trading_assistant/trading_assistant_api/services/user/main.go

package main

import (
	"context"
	"fmt"
	"io"
	"encoding/json"
	"errors"
	"strings"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"database/sql"
	"common/db"
	"common/logger"
	"common/utils"
	"common/types"
	"golang.org/x/crypto/bcrypt"
	"github.com/jackc/pgconn"
)

var pg *sql.DB

func main() {
	port := utils.GetEnv("PORT", "8080")
	srv := &http.Server{
		Addr:    ":" + port,
		Handler: routes(),
	}

	var err error
	pg, err = db.InitPostgres()
	if err != nil {
		logger.L().Fatalf("postgres init: %v", err)
	}
	if utils.GetEnv("MIGRATE_ON_START", "0") == "1" {
		if err := applySchema("db/schema.sql"); err != nil {
			logger.L().Fatalf("apply schema: %v", err)
		}
	}

	logger.L().Printf("user service starting on :%s", port)

	go func() {
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logger.L().Fatalf("listen: %v", err)
		}
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	logger.L().Printf("user service shutting down...")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := srv.Shutdown(ctx); err != nil {
		logger.L().Printf("server shutdown: %v", err)
	}
	_ = db.ClosePostgres()
	logger.L().Printf("user service exited")
}

func routes() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})
	mux.HandleFunc("/version", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("user-service v0.1.0"))
	})
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "hello from user-service")
	})
	mux.HandleFunc("/register", registerHandler)
	mux.HandleFunc("/login", loginHandler)
	return mux
}

type registerReq struct {
	Account  string `json:"account"`
	Password string `json:"password"`
}
type loginReq struct {
	Account  string `json:"account"`
	Password string `json:"password"`
}
type loginResp struct {
	UserID string `json:"user_id"`
}

func registerHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeJSON(w, http.StatusMethodNotAllowed, false, "method not allowed")
		return
	}
	var req registerReq
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeJSON(w, http.StatusBadRequest, false, "invalid json")
		return
	}
	if !validAccount(req.Account) || !validPassword(req.Password) {
		writeJSON(w, http.StatusBadRequest, false, "invalid account or password")
		return
	}
	hashed, err := bcrypt.GenerateFromPassword([]byte(req.Password), 12)
	if err != nil {
		writeJSON(w, http.StatusInternalServerError, false, "internal error")
		return
	}
	tx, err := pg.Begin()
	if err != nil {
		writeJSON(w, http.StatusInternalServerError, false, "internal error")
		return
	}
	defer func() {
		_ = tx.Rollback()
	}()
	var userID string
	if err := tx.QueryRow(`INSERT INTO users DEFAULT VALUES RETURNING user_id`).Scan(&userID); err != nil {
		writeJSON(w, http.StatusInternalServerError, false, "internal error")
		return
	}
	if _, err := tx.Exec(`INSERT INTO user_login_accounts (user_id, value, deleted) VALUES ($1, $2, false)`, userID, req.Account); err != nil {
		if isUniqueViolation(err) {
			writeJSON(w, http.StatusConflict, false, "account exists")
			return
		}
		writeJSON(w, http.StatusInternalServerError, false, "internal error")
		return
	}
	if _, err := tx.Exec(`INSERT INTO user_login_passwords (user_id, value, deleted) VALUES ($1, $2, false)`, userID, string(hashed)); err != nil {
		writeJSON(w, http.StatusInternalServerError, false, "internal error")
		return
	}
	if err := tx.Commit(); err != nil {
		writeJSON(w, http.StatusInternalServerError, false, "internal error")
		return
	}
	writeJSON(w, http.StatusCreated, true, userID)
}

func loginHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeJSON(w, http.StatusMethodNotAllowed, false, "method not allowed")
		return
	}
	var req loginReq
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeJSON(w, http.StatusBadRequest, false, "invalid json")
		return
	}
	if !validAccount(req.Account) || !validPassword(req.Password) {
		writeJSON(w, http.StatusBadRequest, false, "invalid account or password")
		return
	}
	var userID string
	if err := pg.QueryRow(`SELECT user_id FROM user_login_accounts WHERE value = $1 AND deleted = false`, req.Account).Scan(&userID); err != nil {
		writeJSON(w, http.StatusUnauthorized, false, "unauthorized")
		return
	}
	var hashed string
	if err := pg.QueryRow(`SELECT value FROM user_login_passwords WHERE user_id = $1 AND deleted = false`, userID).Scan(&hashed); err != nil {
		writeJSON(w, http.StatusUnauthorized, false, "unauthorized")
		return
	}
	if bcrypt.CompareHashAndPassword([]byte(hashed), []byte(req.Password)) != nil {
		writeJSON(w, http.StatusUnauthorized, false, "unauthorized")
		return
	}
	writeJSON(w, http.StatusOK, true, userID)
}

func validAccount(a string) bool {
	n := len(a)
	return n >= 3 && n <= 100
}
func validPassword(p string) bool {
	n := len(p)
	return n >= 8 && n <= 128
}

func isUniqueViolation(err error) bool {
	var pe *pgconn.PgError
	if errors.As(err, &pe) {
		return pe.Code == "23505"
	}
	if err != nil && (strings.Contains(err.Error(), "duplicate key") || strings.Contains(err.Error(), "unique constraint")) {
		return true
	}
	return false
}

func applySchema(path string) error {
	b, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	_, err = pg.Exec(string(b))
	return err
}

func writeJSON(w http.ResponseWriter, code int, status bool, msg string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(types.Response{Status: status, Message: msg})
}