digital-psychology/apps/api/internal/middleware/device.go

package middleware

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/yuxingu/digital-psychology/apps/api/pkg/response"
)

type ctxKey string

const UserIDKey ctxKey = "user_id"
const DeviceKeyHeader = "X-Device-Key"

// DeviceAuth resolves or creates a Visitor→User via device key.
// If Authorization Bearer session is valid, that account user wins and device rebinds.
func DeviceAuth(pool *pgxpool.Pool) gin.HandlerFunc {
	return func(c *gin.Context) {
		key := c.GetHeader(DeviceKeyHeader)
		if key == "" {
			key = newDeviceKey()
			c.Header(DeviceKeyHeader, key)
			c.Request.Header.Set(DeviceKeyHeader, key)
		}
		if tok := bearerFromHeader(c.GetHeader("Authorization")); tok != "" {
			var uid uuid.UUID
			err := pool.QueryRow(c.Request.Context(), `
				SELECT user_id FROM user_sessions
				WHERE token=$1 AND revoked_at IS NULL AND expires_at > now()`, tok,
			).Scan(&uid)
			if err == nil {
				if banned, berr := userIsBanned(c.Request.Context(), pool, uid); berr != nil {
					response.Fail(c, http.StatusInternalServerError, 50001, "identity unavailable")
					c.Abort()
					return
				} else if banned {
					response.Fail(c, http.StatusForbidden, 40310, "账号已封禁")
					c.Abort()
					return
				}
				_, _ = pool.Exec(c.Request.Context(), `
					INSERT INTO device_identities(device_key, user_id)
					VALUES ($1,$2)
					ON CONFLICT (device_key) DO UPDATE SET user_id=$2, updated_at=now(), deleted_at=NULL`,
					key, uid,
				)
				c.Set(string(UserIDKey), uid.String())
				c.Header(DeviceKeyHeader, key)
				c.Next()
				return
			}
		}
		userID, err := ensureUser(c.Request.Context(), pool, key)
		if err != nil {
			response.Fail(c, http.StatusInternalServerError, 50001, "identity unavailable")
			c.Abort()
			return
		}
		if banned, berr := userIsBanned(c.Request.Context(), pool, userID); berr != nil {
			response.Fail(c, http.StatusInternalServerError, 50001, "identity unavailable")
			c.Abort()
			return
		} else if banned {
			response.Fail(c, http.StatusForbidden, 40310, "账号已封禁")
			c.Abort()
			return
		}
		c.Set(string(UserIDKey), userID.String())
		c.Header(DeviceKeyHeader, key)
		c.Next()
	}
}

func bearerFromHeader(h string) string {
	if len(h) < 8 {
		return ""
	}
	if strings.EqualFold(h[:7], "bearer ") {
		return strings.TrimSpace(h[7:])
	}
	return ""
}

// RequireRegistered rejects anonymous (no phone) users.
func RequireRegistered(pool *pgxpool.Pool) gin.HandlerFunc {
	return func(c *gin.Context) {
		userID, ok := UserIDFromContext(c)
		if !ok {
			response.Fail(c, http.StatusUnauthorized, 40100, "unauthorized")
			c.Abort()
			return
		}
		var okReg bool
		err := pool.QueryRow(c.Request.Context(), `
			SELECT EXISTS(
				SELECT 1 FROM users WHERE id=$1 AND phone IS NOT NULL AND deleted_at IS NULL
			)`, userID).Scan(&okReg)
		if err != nil || !okReg {
			response.Fail(c, http.StatusUnauthorized, 40112, "请先登录后再使用")
			c.Abort()
			return
		}
		c.Next()
	}
}

// UserIDFromContext returns the authenticated user id.
func UserIDFromContext(c *gin.Context) (uuid.UUID, bool) {
	v, ok := c.Get(string(UserIDKey))
	if !ok {
		return uuid.Nil, false
	}
	id, err := uuid.Parse(v.(string))
	return id, err == nil
}

func ensureUser(ctx context.Context, pool *pgxpool.Pool, deviceKey string) (uuid.UUID, error) {
	var userID *uuid.UUID
	err := pool.QueryRow(ctx, `
		SELECT user_id FROM device_identities
		WHERE device_key=$1 AND deleted_at IS NULL`, deviceKey,
	).Scan(&userID)
	if err == nil && userID != nil {
		return *userID, nil
	}
	if err != nil && !errors.Is(err, pgx.ErrNoRows) {
		return uuid.Nil, err
	}

	tx, err := pool.Begin(ctx)
	if err != nil {
		return uuid.Nil, err
	}
	defer tx.Rollback(ctx)

	var uid uuid.UUID
	if err := tx.QueryRow(ctx,
		`INSERT INTO users DEFAULT VALUES RETURNING id`,
	).Scan(&uid); err != nil {
		return uuid.Nil, err
	}
	if _, err := tx.Exec(ctx, `
		INSERT INTO device_identities(device_key, user_id)
		VALUES ($1,$2)
		ON CONFLICT (device_key) DO UPDATE SET user_id=EXCLUDED.user_id, updated_at=now()`,
		deviceKey, uid,
	); err != nil {
		return uuid.Nil, err
	}
	if err := tx.Commit(ctx); err != nil {
		return uuid.Nil, err
	}
	return uid, nil
}

func newDeviceKey() string {
	b := make([]byte, 16)
	_, _ = rand.Read(b)
	return "dev_" + hex.EncodeToString(b)
}

func userIsBanned(ctx context.Context, pool *pgxpool.Pool, userID uuid.UUID) (bool, error) {
	var status string
	err := pool.QueryRow(ctx, `
		SELECT status FROM users WHERE id=$1 AND deleted_at IS NULL`, userID).Scan(&status)
	if errors.Is(err, pgx.ErrNoRows) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return status == "banned", nil
}