auth/server/server.go

package server

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"time"

	"git.t-juice.club/microfilm/auth"
	"git.t-juice.club/microfilm/auth/store"
	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
	"github.com/nats-io/nats.go"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.21.0"
	"go.opentelemetry.io/otel/trace"
)

const DefaultTokenDuration time.Duration = 24 * time.Hour

type Server struct {
	Logger *slog.Logger

	http.Server
	store      store.AuthStore
	config     *Config
	nats       *nats.EncodedConn
	userClient *UserClient

	signingKey *ecdsa.PrivateKey
}

func NewServer(config *Config) (*Server, error) {
	srv := &Server{}

	tp, err := tracerProvider("jaeger:4318")
	if err != nil {
		return nil, err
	}
	otel.SetTracerProvider(tp)

	r := chi.NewRouter()

	r.Use(middleware.RequestID)
	r.Use(srv.MiddlewareLogging)
	r.Use(srv.MiddlewareTracing)

	r.Get("/key", srv.PubkeyHandler)
	r.Post("/{id}/token", srv.TokenHandler)

	srv.Handler = r
	srv.Addr = config.ListenAddr

	srv.config = config

	srv.Logger = slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{
		Level: slog.LevelDebug,
	}))

	srv.store = store.NewMemoryAuthStore()

	conn, err := nats.Connect(config.NATSAddr)
	if err != nil {
		return nil, err
	}
	encoded, err := nats.NewEncodedConn(conn, "json")
	if err != nil {
		return nil, err
	}

	srv.nats = encoded
	srv.userClient = NewUserClient(config.UserServiceBaseURL)

	// Generate keys
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}
	srv.signingKey = privateKey

	return srv, nil
}

func tracerProvider(url string) (*sdktrace.TracerProvider, error) {
	exp, err := otlptracehttp.New(context.Background(), otlptracehttp.WithEndpoint(url), otlptracehttp.WithInsecure())
	if err != nil {
		return nil, err
	}

	res := resource.NewWithAttributes(semconv.SchemaURL,
		semconv.ServiceName("mf-auth"),
		semconv.ServiceVersion(auth.Version),
	)
	tp := sdktrace.NewTracerProvider(
		sdktrace.WithBatcher(exp, sdktrace.WithBatchTimeout(time.Second)),
		sdktrace.WithResource(res),
	)
	otel.SetTracerProvider(tp)
	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{}))

	return tp, nil
}

func InfoHandler(w http.ResponseWriter, r *http.Request) {
	enc := json.NewEncoder(w)

	data := &auth.InfoResponse{
		Version: auth.Version,
	}

	_ = enc.Encode(data)
}

func WriteError(w http.ResponseWriter, response auth.ErrorResponse) {
	encoder := json.NewEncoder(w)
	w.WriteHeader(response.Status)
	_ = encoder.Encode(&response)
}

func (s *Server) PubkeyHandler(w http.ResponseWriter, r *http.Request) {
	span := trace.SpanFromContext(r.Context())

	span.AddEvent("Start marshalling public key.")
	enc := json.NewEncoder(w)
	key, err := x509.MarshalPKIXPublicKey(s.signingKey.Public())
	if err != nil {
		s.Logger.Error("Unable to marshal public key.", "error", err)
		WriteError(w, auth.ErrorResponse{
			Status:  http.StatusInternalServerError,
			Message: "Unable to marshal public key",
		})
		return
	}
	span.AddEvent("Finished marshalling public key.")
	response := auth.PubkeyResponse{
		PubKey: key,
	}

	_ = enc.Encode(&response)
}

func (s *Server) TokenHandler(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	decoder := json.NewDecoder(r.Body)
	defer r.Body.Close()

	userIdentifier := chi.URLParam(r, "id")
	if userIdentifier == "" {
		WriteError(w, auth.ErrorResponse{
			Status:  http.StatusBadRequest,
			Message: "Invalid user identifier.",
		})
		return
	}

	var request auth.TokenRequest
	if err := decoder.Decode(&request); err != nil {
		WriteError(w, auth.ErrorResponse{
			Status:  http.StatusBadRequest,
			Message: fmt.Sprintf("Error parsing request: %s", err),
		})
		return
	}

	if err := s.userClient.VerifyUserPassword(ctx, userIdentifier, request.Password); err != nil {
		WriteError(w, auth.ErrorResponse{
			Status:  http.StatusUnauthorized,
			Message: fmt.Sprintf("Unable to verify password: %s", err),
		})
		return
	}

	u, err := s.userClient.GetUser(ctx, userIdentifier)
	if err != nil {
		WriteError(w, auth.ErrorResponse{
			Status:  http.StatusUnauthorized,
			Message: fmt.Sprintf("Unable to get user details: %s", err),
		})
		return
	}

	exp := time.Now().Add(DefaultTokenDuration)
	claims := auth.MicrofilmClaims{
		Role: u.Role,
		RegisteredClaims: jwt.RegisteredClaims{
			Issuer:    "microfilm",
			Subject:   userIdentifier,
			Audience:  []string{"microfilm"},
			ExpiresAt: jwt.NewNumericDate(exp),
			NotBefore: jwt.NewNumericDate(time.Now()),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
			ID:        uuid.Must(uuid.NewRandom()).String(),
		},
	}
	token := jwt.NewWithClaims(jwt.SigningMethodES256, claims)
	tokenString, err := token.SignedString(s.signingKey)
	if err != nil {
		WriteError(w, auth.ErrorResponse{
			Status:  http.StatusInternalServerError,
			Message: fmt.Sprintf("Unable to sign token: %s", err),
		})
		return
	}

	if err := s.store.Add(store.RevokableToken{
		ID:        claims.RegisteredClaims.ID,
		Subject:   claims.RegisteredClaims.Subject,
		ExpiresAt: exp,
	}); err != nil {
		WriteError(w, auth.ErrorResponse{
			Status:  http.StatusInternalServerError,
			Message: fmt.Sprintf("Unable to store token for revocation: %s", err),
		})
		return
	}

	// Publish message
	subject := fmt.Sprintf("%s.%s", s.config.BaseSubject, "issued")
	msg := auth.MsgTokenIssued{
		Message: "Token issued.",
		ID:      claims.RegisteredClaims.ID,
		Subject: claims.RegisteredClaims.Subject,
	}
	if err := s.nats.Publish(subject, msg); err != nil {
		s.Logger.Error("Unable to publish message.", "error", err)
		WriteError(w, auth.ErrorResponse{
			Status:  http.StatusInternalServerError,
			Message: fmt.Sprintf("Unable to store token for revocation: %s", err),
		})
		return
	}

	response := auth.TokenResponse{
		Token: tokenString,
	}

	w.WriteHeader(http.StatusOK)
	encoder := json.NewEncoder(w)
	_ = encoder.Encode(&response)
}