alpha

handlers/handler.go

@@ -1,9 +1,27 @@
 package handlers
 
 import (
+	"encoding/json"
+	"net/http"
+
 	"github.com/jmoiron/sqlx"
 )
 
 type Handler struct {
 	DB *sqlx.DB
 }
+
+type Response struct {
+	Code    int         `json:"code"`
+	Message string      `json:"message"`
+	Data    interface{} `json:"data,omitempty"`
+}
+
+func WriteJSON(w http.ResponseWriter, data interface{}, code int) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(code)
+	json.NewEncoder(w).Encode(data)
+}
+func WriteError(w http.ResponseWriter, message string, code int) {
+	WriteJSON(w, Response{Code: code, Message: message}, code)
+}

handlers/login.go

@@ -5,11 +5,18 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"time"
 
+	"github.com/golang-jwt/jwt"
 	"github.com/spf13/viper"
 )
 
-var code2Session = "https://api.weixin.qq.com/sns/jscode2session?appid=%s&secret=%s&js_code=%s&grant_type=authorization_code"
+var wxClient = &http.Client{
+	Timeout: 10 * time.Second,
+	Transport: &http.Transport{
+		MaxIdleConnsPerHost: 10,
+	},
+}
 
 type LoginRequest struct {
 	Code string `json:"code"`
@@ -27,8 +34,8 @@ type LoginResponse struct {
 func getLoginResponse(code string) (*LoginResponse, error) {
 	appid := viper.GetString("wechat.appid")
 	secret := viper.GetString("wechat.secret")
-	url := fmt.Sprintf(code2Session, appid, secret, code)
-	resp, err := http.Get(url)
+	url := fmt.Sprintf("https://api.weixin.qq.com/sns/jscode2session?appid=%s&secret=%s&js_code=%s&grant_type=authorization_code", appid, secret, code)
+	resp, err := wxClient.Get(url)
 	if err != nil {
 		return nil, err
 	}
@@ -45,49 +52,71 @@ func getLoginResponse(code string) (*LoginResponse, error) {
 	return &loginResponse, nil
 }
 
-func (h *Handler) Login(w http.ResponseWriter, r *http.Request) {
+func generateJWT(openid string) (string, error) {
+	tokenTTL := viper.GetDuration("auth.ttl")
+	secret := viper.GetString("auth.secret")
+	if secret == "" {
+		return "", fmt.Errorf("auth secret not set")
+	}
 
+	claims := jwt.MapClaims{
+		"openid": openid,
+		"exp":    time.Now().Add(tokenTTL).Unix(),
+		"iat":    time.Now().Unix(),
+		"iss":    viper.GetString("server.name"),
+	}
+
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
+	signedToken, err := token.SignedString([]byte(secret))
+	if err != nil {
+		return "", err
+	}
+	return signedToken, nil
+}
+
+func (h *Handler) Login(w http.ResponseWriter, r *http.Request) {
 	var req LoginRequest
 	body, err := io.ReadAll(r.Body)
 	if err != nil {
-		http.Error(w, "Failed to read request body", http.StatusBadRequest)
+		WriteError(w, "Failed to read request body", http.StatusBadRequest)
 		return
 	}
 	if err := json.Unmarshal(body, &req); err != nil {
-		http.Error(w, "Failed to parse request body", http.StatusBadRequest)
+		WriteError(w, "Failed to parse request body", http.StatusBadRequest)
 		return
 	}
 
 	loginResponse, err := getLoginResponse(req.Code)
 	if err != nil {
-		http.Error(w, "Failed to get session key", http.StatusInternalServerError)
+		WriteError(w, "Failed to get login response", http.StatusInternalServerError)
 		return
 	}
 
-	// // 插入或更新用户的openid和sessionid
-	// query := `
-	// 	INSERT INTO users (openid, sessionid)
-	// 	VALUES ($1, $2)
-	// 	ON CONFLICT (openid) DO UPDATE SET sessionid = $2
-	// 	RETURNING id;
-	// `
+	// 插入或更新用户的openid和session_key
+	query := `
+		INSERT INTO users (openid, session_key)
+		VALUES ($1, $2)
+		ON CONFLICT (openid) DO UPDATE SET session_key = $2
+		RETURNING id;
+	`
 
-	// var ID int
-	// if err := h.DB.QueryRow(query, loginResponse.OpenId, loginResponse.SessionKey).Scan(&ID); err != nil {
-	// 	http.Error(w, "Failed to log in user", http.StatusInternalServerError)
-	// 	return
-	// }
+	var ID int
+	if err := h.DB.QueryRow(query, loginResponse.OpenId, loginResponse.SessionKey).Scan(&ID); err != nil {
+		WriteError(w, "Failed to log in user", http.StatusInternalServerError)
+		return
+	}
+
+	token, err := generateJWT(loginResponse.OpenId)
+	if err != nil {
+		WriteError(w, "Failed to generate JWT token", http.StatusInternalServerError)
+		return
+	}
 
 	data := map[string]interface{}{
+		"token":       token,
 		"openid":      loginResponse.OpenId,
 		"session_key": loginResponse.SessionKey,
 	}
-	respData, err := json.Marshal(data)
-	if err != nil {
-		http.Error(w, "Failed to marshal response data", http.StatusInternalServerError)
-		return
-	}
-	w.Header().Set("Content-Type", "application/json")
-	w.WriteHeader(http.StatusOK)
-	w.Write([]byte(respData))
+
+	WriteJSON(w, Response{Code: 0, Message: "Success", Data: data}, http.StatusOK)
 }

handlers/otp.go

@@ -2,6 +2,7 @@ package handlers
 
 import (
 	"encoding/json"
+	"log"
 	"net/http"
 )
 
@@ -19,43 +20,53 @@ type OTP struct {
 func (h *Handler) UpdateOrCreateOtp(w http.ResponseWriter, r *http.Request) {
 	var req OtpRequest
 	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
-		http.Error(w, "Invalid request payload", http.StatusBadRequest)
+		WriteError(w, "Failed to parse request body", http.StatusBadRequest)
 		return
 	}
 
-	num := len(*req.Token)
+	if req.OpenID == "" {
+		WriteError(w, "OpenID is required", http.StatusBadRequest)
+		return
+	}
+
+	if req.Token == nil || len(*req.Token) == 0 {
+		WriteError(w, "Token is required", http.StatusBadRequest)
+		return
+	}
+
+	log.Printf("Saving OTP for user: %s token count:: %d", req.OpenID, len(*req.Token))
 
 	// 插入或更新 OTP 记录
 	query := `
 		INSERT INTO otp (openid, num, token)
 		VALUES ($1, $2, $3)
+		ON CONFLICT (openid) DO UPDATE SET num = EXCLUDED.num, token = EXCLUDED.token
+	}
 	`
 
-	_, err := h.DB.Exec(query, req.OpenID, req.Token, num)
+	_, err := h.DB.Exec(query, req.OpenID, len(*req.Token), req.Token)
 	if err != nil {
-		http.Error(w, "Failed to update or create OTP", http.StatusInternalServerError)
+		WriteError(w, "Failed to update or create OTP", http.StatusInternalServerError)
 		return
 	}
 
-	w.WriteHeader(http.StatusOK)
-	w.Write([]byte("OTP updated or created successfully"))
+	WriteJSON(w, Response{Code: 0, Message: "Success"}, http.StatusOK)
 }
 
 func (h *Handler) GetOtp(w http.ResponseWriter, r *http.Request) {
 	openid := r.URL.Query().Get("openid")
 	if openid == "" {
-		http.Error(w, "未登录", http.StatusBadRequest)
+		WriteError(w, "OpenID is required", http.StatusBadRequest)
 		return
 	}
 
 	var otp OtpRequest
 
-	err := h.DB.Get(&otp, "SELECT token, num, openid FROM otp WHERE openid=$1", openid)
+	err := h.DB.Get(&otp, "SELECT openid, token, num FROM otp WHERE openid=$1", openid)
 	if err != nil {
-		http.Error(w, "OTP not found", http.StatusNotFound)
+		WriteError(w, "Failed to get OTP", http.StatusInternalServerError)
 		return
 	}
 
-	w.Header().Set("Content-Type", "application/json")
-	json.NewEncoder(w).Encode(otp)
+	WriteJSON(w, Response{Code: 0, Message: "Success", Data: otp}, http.StatusOK)
 }