// Copyright (c) 2026 Charles KWON OhJun (charleskwonohjun@gmail.com)
// All rights reserved.

// auth.go — password / md5 authentication for the pgserver.
//
// Roles + credentials live in an in-memory registry managed via
// `PG_ADD_ROLE(name, password)` HB_FUNC (see register.go). At
// startup the bootstrap PRG calls PG_ADD_ROLE for every account
// that should be allowed in; `trust` mode bypasses lookup
// entirely so single-user / dev setups don't need credentials.
//
// SCRAM-SHA-256 is Phase 5.1 — pgx falls back to MD5 cleanly
// when the server advertises only md5, so v1.0 functional
// coverage is complete with the two methods here.

package pgserver

import (
	"crypto/md5"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"sync"

	"github.com/jackc/pgx/v5/pgproto3"
)

// role captures a stored credential. PasswordPlain is held so
// cleartext-password mode (the simplest path) doesn't need a
// separate verification table. MD5 mode computes the canonical
// hash from PasswordPlain at challenge time — matches what
// Postgres does internally when md5 is configured against a
// scram-stored password before users opt in to SCRAM.
type role struct {
	Name          string
	PasswordPlain string
}

var (
	roleMu  sync.RWMutex
	roleMap = map[string]*role{}
)

// AddRole registers a user + password. Replaces any prior entry
// with the same name (so a bootstrap PRG can re-add roles on
// restart without first DROPping them).
func AddRole(name, password string) {
	roleMu.Lock()
	defer roleMu.Unlock()
	roleMap[name] = &role{Name: name, PasswordPlain: password}
}

// RemoveRole drops a registered user. No-op if unknown.
func RemoveRole(name string) {
	roleMu.Lock()
	defer roleMu.Unlock()
	delete(roleMap, name)
}

// lookupRole resolves a role by name. Returns nil if absent.
func lookupRole(name string) *role {
	roleMu.RLock()
	defer roleMu.RUnlock()
	return roleMap[name]
}

// authenticate runs the auth handshake based on the server's
// configured AuthMode. The client identity (s.user) has already
// been recorded from the StartupMessage; we look it up in the
// role registry and execute the appropriate challenge.
func (s *session) authenticate() error {
	switch s.srv.cfg.AuthMode {
	case "", "trust":
		s.send(&pgproto3.AuthenticationOk{})
		return nil
	case "password":
		return s.authPassword()
	case "md5":
		return s.authMD5()
	default:
		s.sendError("28000",
			"auth mode "+s.srv.cfg.AuthMode+" not implemented (use trust/password/md5)")
		return errAuthRejected
	}
}

// authPassword does the cleartext-password exchange. The wire
// payload is plaintext, so this is intended for TLS-protected
// links only — emit a warning if the connection isn't tls.Server
// (deferred until Phase 6 wires up TLS detection on session).
func (s *session) authPassword() error {
	s.send(&pgproto3.AuthenticationCleartextPassword{})
	msg, err := s.be.Receive()
	if err != nil {
		return err
	}
	pwd, ok := msg.(*pgproto3.PasswordMessage)
	if !ok {
		s.sendError("28000", "expected PasswordMessage")
		return errAuthRejected
	}
	r := lookupRole(s.user)
	if r == nil || r.PasswordPlain != pwd.Password {
		s.sendError("28P01",
			fmt.Sprintf("password authentication failed for user %q", s.user))
		return errAuthRejected
	}
	s.send(&pgproto3.AuthenticationOk{})
	return nil
}

// authMD5 implements the libpq MD5 password challenge:
//
//	server sends: AuthenticationMD5Password{Salt: 4 random bytes}
//	client returns: "md5" || md5_hex( md5_hex(password || user) || salt )
//	server verifies by recomputing with the stored plaintext
//
// MD5 is no longer recommended (libpq 14+ default is SCRAM), but
// every PG client implements it as a fallback. Adequate for v1.0
// over loopback or a trusted network; Phase 5.1 lands SCRAM.
func (s *session) authMD5() error {
	var salt [4]byte
	if _, err := rand.Read(salt[:]); err != nil {
		s.sendError("XX000", "auth: rng failure")
		return err
	}
	s.send(&pgproto3.AuthenticationMD5Password{Salt: salt})
	msg, err := s.be.Receive()
	if err != nil {
		return err
	}
	pwd, ok := msg.(*pgproto3.PasswordMessage)
	if !ok {
		s.sendError("28000", "expected PasswordMessage")
		return errAuthRejected
	}
	r := lookupRole(s.user)
	if r == nil {
		s.sendError("28P01",
			fmt.Sprintf("md5 authentication failed for user %q", s.user))
		return errAuthRejected
	}
	expected := md5Challenge(r.PasswordPlain, s.user, salt[:])
	if pwd.Password != expected {
		s.sendError("28P01",
			fmt.Sprintf("md5 authentication failed for user %q", s.user))
		return errAuthRejected
	}
	s.send(&pgproto3.AuthenticationOk{})
	return nil
}

// md5Challenge reproduces libpq's md5 client computation so we
// can compare against the value the client sent.
func md5Challenge(password, user string, salt []byte) string {
	inner := md5.Sum([]byte(password + user))
	innerHex := hex.EncodeToString(inner[:])
	outer := md5.Sum(append([]byte(innerHex), salt...))
	return "md5" + hex.EncodeToString(outer[:])
}

// sentinelError lets the run() loop bail out without typing the
// "fmt.Errorf" boilerplate at every call site.
type sentinelError string

func (e sentinelError) Error() string { return string(e) }

const errAuthRejected = sentinelError("pgserver: authentication rejected")