byteclaw/internal/server/auth.go

package server

import (
	"context"
	"crypto/subtle"
	"errors"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"golang.org/x/crypto/bcrypt"

	"github.com/crypto/byteclaw/internal/store"
)

// Principal is the authenticated caller. Admin (static auth_token) sees all sessions.
type Principal struct {
	UserID   int64
	Username string
	Admin    bool
	token    string
}

// Scope is the user filter for store queries.
func (p Principal) Scope() int64 {
	if p.Admin {
		return store.AllUsers
	}
	return p.UserID
}

// OwnerID is the user_id new sessions get.
func (p Principal) OwnerID() int64 {
	if p.Admin {
		return 0
	}
	return p.UserID
}

type ctxKey struct{}

func principal(ctx context.Context) Principal {
	p, _ := ctx.Value(ctxKey{}).(Principal)
	return p
}

const (
	defaultTokenTTL = 30 * 24 * time.Hour
	loginWindow     = 15 * time.Minute
	loginMaxFails   = 10
)

// dummyHash keeps the timing of "unknown user" close to "wrong password".
var dummyHash, _ = bcrypt.GenerateFromPassword([]byte("byteclaw-dummy-password"), bcrypt.DefaultCost)

// limiter counts failed attempts per IP in a fixed window.
type limiter struct {
	mu   sync.Mutex
	hits map[string][]time.Time
	max  int
	win  time.Duration
	now  func() time.Time
}

func newLimiter(max int, win time.Duration) *limiter {
	return &limiter{hits: map[string][]time.Time{}, max: max, win: win, now: time.Now}
}

func (l *limiter) prune(key string) []time.Time {
	cut := l.now().Add(-l.win)
	h := l.hits[key]
	i := 0
	for i < len(h) && h[i].Before(cut) {
		i++
	}
	h = h[i:]
	if len(h) == 0 {
		delete(l.hits, key)
	} else {
		l.hits[key] = h
	}
	return h
}

func (l *limiter) Blocked(key string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	return len(l.prune(key)) >= l.max
}

func (l *limiter) Fail(key string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.hits[key] = append(l.prune(key), l.now())
}

func (l *limiter) Reset(key string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	delete(l.hits, key)
}

func clientIP(r *http.Request) string {
	h, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return h
}

func (s *Server) initAuth() {
	s.authOnce.Do(func() {
		if s.limiter == nil {
			s.limiter = newLimiter(loginMaxFails, loginWindow)
		}
		if s.TokenTTL <= 0 {
			s.TokenTTL = defaultTokenTTL
		}
	})
}

func bearer(r *http.Request) string {
	if t, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer "); ok {
		return strings.TrimSpace(t)
	}
	// Browsers cannot set headers on WebSocket upgrades: allow ?token= there only.
	if r.URL.Path == "/api/ws" {
		return r.URL.Query().Get("token")
	}
	return ""
}

// auth accepts the static admin token (constant-time) or a user token.
func (s *Server) auth(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got := bearer(r)
		var p Principal
		switch {
		case got == "":
			writeErr(w, http.StatusUnauthorized, "unauthorized")
			return
		case s.Token != "" && subtle.ConstantTimeCompare([]byte(got), []byte(s.Token)) == 1:
			p = Principal{Admin: true, Username: "admin"}
		default:
			u, err := s.Store.UserByToken(r.Context(), got)
			if err != nil {
				writeErr(w, http.StatusUnauthorized, "unauthorized")
				return
			}
			p = Principal{UserID: u.ID, Username: u.Username}
		}
		p.token = got
		h.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), ctxKey{}, p)))
	})
}

type credentials struct {
	Username string `json:"username"`
	Password string `json:"password"`
}

func validUsername(u string) bool {
	if n := utf8.RuneCountInString(u); n < 3 || n > 32 {
		return false
	}
	for _, c := range u {
		if !(c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || c == '_' || c == '-' || c == '.') {
			return false
		}
	}
	return true
}

func validPassword(p string) bool {
	// bcrypt only uses the first 72 bytes; reject longer input instead of silently truncating.
	return utf8.RuneCountInString(p) >= 8 && len(p) <= 72
}

func (s *Server) issue(w http.ResponseWriter, r *http.Request, u store.User, code int) {
	tok, exp, err := s.Store.IssueToken(r.Context(), u.ID, s.TokenTTL)
	if err != nil {
		s.storeErr(w, err)
		return
	}
	writeJSON(w, code, map[string]any{"token": tok, "expires_at": exp, "user": map[string]any{"id": u.ID, "username": u.Username}})
}

func (s *Server) register(w http.ResponseWriter, r *http.Request) {
	s.initAuth()
	if !s.AllowRegistration {
		writeErr(w, http.StatusForbidden, "registration disabled")
		return
	}
	ip := clientIP(r)
	if s.limiter.Blocked(ip) {
		writeErr(w, http.StatusTooManyRequests, "too many attempts, try again later")
		return
	}
	var c credentials
	if !decode(w, r, &c) {
		return
	}
	c.Username = strings.TrimSpace(c.Username)
	if !validUsername(c.Username) {
		writeErr(w, http.StatusBadRequest, "username must be 3-32 characters (letters, digits, . _ -)")
		return
	}
	if !validPassword(c.Password) {
		writeErr(w, http.StatusBadRequest, "password must be 8-72 characters")
		return
	}
	s.limiter.Fail(ip) // every registration counts against the IP budget
	hash, err := bcrypt.GenerateFromPassword([]byte(c.Password), bcrypt.DefaultCost)
	if err != nil {
		s.storeErr(w, err)
		return
	}
	u, err := s.Store.CreateUser(r.Context(), c.Username, string(hash))
	if errors.Is(err, store.ErrUserExists) {
		writeErr(w, http.StatusConflict, "username already taken")
		return
	}
	if err != nil {
		s.storeErr(w, err)
		return
	}
	s.issue(w, r, u, http.StatusCreated)
}

func (s *Server) login(w http.ResponseWriter, r *http.Request) {
	s.initAuth()
	ip := clientIP(r)
	if s.limiter.Blocked(ip) {
		writeErr(w, http.StatusTooManyRequests, "too many attempts, try again later")
		return
	}
	var c credentials
	if !decode(w, r, &c) {
		return
	}
	u, err := s.Store.UserByName(r.Context(), strings.TrimSpace(c.Username))
	hash := dummyHash
	if err == nil {
		hash = []byte(u.PasswordHash())
	} else if !errors.Is(err, store.ErrNotFound) {
		s.storeErr(w, err)
		return
	}
	if bcrypt.CompareHashAndPassword(hash, []byte(c.Password)) != nil || err != nil {
		s.limiter.Fail(ip)
		writeErr(w, http.StatusUnauthorized, "invalid username or password")
		return
	}
	s.limiter.Reset(ip)
	s.issue(w, r, u, http.StatusOK)
}

func (s *Server) logout(w http.ResponseWriter, r *http.Request) {
	p := principal(r.Context())
	if !p.Admin {
		if err := s.Store.RevokeToken(r.Context(), p.token); err != nil {
			s.storeErr(w, err)
			return
		}
	}
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) me(w http.ResponseWriter, r *http.Request) {
	p := principal(r.Context())
	writeJSON(w, 200, map[string]any{"id": p.UserID, "username": p.Username, "admin": p.Admin})
}

// owns writes 404 and returns false unless the caller may access session id.
func (s *Server) owns(w http.ResponseWriter, r *http.Request, id string) bool {
	if err := s.Store.SessionOwnedBy(r.Context(), principal(r.Context()).Scope(), id); err != nil {
		s.storeErr(w, err)
		return false
	}
	return true
}