package main

// LocalSend protocol v2.2, receive side.
//
// SnapBeam speaks the protocol documented at https://github.com/localsend/protocol
// so that the official LocalSend app — a separate Apache-2.0 project by its own
// authors — can discover this computer and send to it. SnapBeam is not
// LocalSend, is not affiliated with it, and ships none of its code; this file
// is an independent implementation written from the published specification.
//
// Implemented here (protocol section numbers refer to that document):
//
//	POST /api/localsend/v2/register        3.1 / 3.2  two-way discovery
//	GET  /api/localsend/v2/info            6.1        device information
//	POST /api/localsend/v2/prepare-upload  4.1        metadata, PIN, tokens
//	POST /api/localsend/v2/upload          4.2        the file bytes
//	POST /api/localsend/v2/cancel          4.3        sender gave up
//
// Not implemented, deliberately and stated in the README: the download API
// (5.x), HTTPS/mTLS transport, and protocol v1. Requests to v1 routes get a
// 404 rather than a wrong answer.

import (
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// protocolVersion is the version SnapBeam announces and implements. The
// official implementation announces "2.2" and accepts "2.0" from peers.
const protocolVersion = "2.2"

// defaultLocalSendPort is both the HTTP port and the multicast port in the
// protocol's default configuration (section 1).
const defaultLocalSendPort = 53317

// multicastGroup is the group the protocol specifies. It sits inside
// 224.0.0.0/24 because some Android devices refuse any other range.
var multicastGroup = net.IPv4(224, 0, 0, 167)

// maxPINAttempts matches the reference implementation: three wrong PINs from
// one address and that address gets 429 until the receiver restarts.
const maxPINAttempts = 3

// maxMetadataBytes caps a prepare-upload body. The metadata for a realistic
// transfer is kilobytes; a megabyte of JSON is an attack, not a photo album.
const maxMetadataBytes = 4 << 20

// sessionIdleTimeout releases the single session slot when a sender vanishes
// mid-transfer. Without it one dead phone would block the receiver with 409
// until it was restarted. This is SnapBeam's own policy; the protocol does not
// specify a timeout.
const sessionIdleTimeout = 5 * time.Minute

// ---------------------------------------------------------------------------
// Wire types (protocol sections 3, 4 and 6)
// ---------------------------------------------------------------------------

// deviceInfo is the device descriptor that appears as the multicast
// announcement body, as the /register request body, and as the "info" object
// inside a prepare-upload request.
type deviceInfo struct {
	Alias       string `json:"alias"`
	Version     string `json:"version"`
	DeviceModel string `json:"deviceModel,omitempty"`
	DeviceType  string `json:"deviceType,omitempty"`
	Fingerprint string `json:"fingerprint"`
	Port        int    `json:"port"`
	Protocol    string `json:"protocol"`
	Download    bool   `json:"download"`
}

// announcement is a deviceInfo plus the "announce" flag that marks a message
// as an announcement rather than a reply.
//
// Ambiguity, and how it is resolved: the published v2.2 document shows
// "announce": true on the announcement and describes an optional UDP reply
// carrying "announce": false, while the current reference implementation
// serialises the flag on the way out but ignores it on the way in and never
// replies over UDP at all. SnapBeam follows the document: it sends
// "announce": true, treats a message with the flag absent as an announcement
// (that is what the reference sends when read strictly), and never replies to
// one carrying "announce": false — which is what stops two SnapBeams from
// answering each other forever.
type announcement struct {
	deviceInfo
	Announce bool `json:"announce"`
}

// registerResponse is the reply to POST /register (section 3.2). It carries no
// port: the peer already knows which socket it connected to.
type registerResponse struct {
	Alias       string `json:"alias"`
	Version     string `json:"version"`
	DeviceModel string `json:"deviceModel,omitempty"`
	DeviceType  string `json:"deviceType,omitempty"`
	Fingerprint string `json:"fingerprint"`
	Download    bool   `json:"download"`
}

// infoResponse is the reply to GET /info (section 6.1). Same shape as
// registerResponse; kept separate because the two routes are free to diverge.
type infoResponse = registerResponse

// fileMetadata carries the original timestamps, added in protocol 2.1.
type fileMetadata struct {
	Modified string `json:"modified,omitempty"`
	Accessed string `json:"accessed,omitempty"`
}

// fileDTO is one offered file (section 4.1).
//
// Ambiguity, and how it is resolved: the published document names the checksum
// field "sha256"; the Dart client that shipped for years encodes it as "hash".
// The current Rust core uses "sha256". SnapBeam accepts either on the way in
// and verifies whichever it was given.
type fileDTO struct {
	ID       string        `json:"id"`
	FileName string        `json:"fileName"`
	Size     int64         `json:"size"`
	FileType string        `json:"fileType"`
	SHA256   string        `json:"sha256,omitempty"`
	Hash     string        `json:"hash,omitempty"`
	Preview  string        `json:"preview,omitempty"`
	Metadata *fileMetadata `json:"metadata,omitempty"`
}

// checksum returns the declared SHA-256 of the file, under whichever of the
// two field names the sender used, or "" when none was declared.
func (f fileDTO) checksum() string {
	if f.SHA256 != "" {
		return f.SHA256
	}
	return f.Hash
}

// prepareUploadRequest is the body of POST /prepare-upload (section 4.1).
type prepareUploadRequest struct {
	Info  deviceInfo         `json:"info"`
	Files map[string]fileDTO `json:"files"`
}

// prepareUploadResponse is its reply: a session id and one token per accepted
// file. Files left out of the map were not accepted.
type prepareUploadResponse struct {
	SessionID string            `json:"sessionId"`
	Files     map[string]string `json:"files"`
}

// ---------------------------------------------------------------------------
// Session state
// ---------------------------------------------------------------------------

type fileStatus int

const (
	filePending fileStatus = iota
	fileInProgress
	fileDone
	fileFailed
)

type sessionFile struct {
	dto    fileDTO
	token  string
	status fileStatus
}

// uploadSession is the one transfer the receiver will accept at a time. The
// protocol requires exactly this: a second sender gets 409 (section 4.1).
type uploadSession struct {
	id       string
	senderIP string
	alias    string
	files    map[string]*sessionFile
	started  time.Time
	touched  time.Time
}

// ---------------------------------------------------------------------------
// The receiver
// ---------------------------------------------------------------------------

// receiver holds everything both the LocalSend routes and the built-in browser
// page need: where files go, how big they may be, and the pairing code that
// guards both doors.
type receiver struct {
	Alias       string
	DeviceModel string
	DeviceType  string
	Fingerprint string
	Port        int
	Dir         string
	MaxBytes    int64
	PairCode    string
	LocalSend   bool

	// report is called once per item that lands, from whichever path
	// delivered it. Both paths print the same line because both call this.
	report func(savedItem)

	// onRegister is called when another device introduces itself over HTTP.
	// This is the half of discovery that does not arrive by multicast: a
	// device that was already running answers OUR announcement by registering
	// with us, and this is where that answer surfaces.
	onRegister func(peer)

	mu       sync.Mutex
	sess     *uploadSession
	pinFails map[string]int
	received []savedItem
}

func newReceiver() *receiver {
	return &receiver{pinFails: map[string]int{}}
}

// newFingerprint generates the random identifier the protocol asks for in
// plain-HTTP mode (section 2). It is regenerated every run, so it identifies a
// running receiver, never the machine across restarts.
func newFingerprint() (string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("cannot generate a device fingerprint: %w", err)
	}
	return hex.EncodeToString(b), nil
}

// newToken generates a session id or a per-file token.
func newToken() (string, error) {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("cannot generate a token: %w", err)
	}
	return hex.EncodeToString(b), nil
}

// newPairCode generates the six-digit code shown next to the QR code. It is
// drawn from crypto/rand with rejection sampling, so all 1,000,000 codes are
// equally likely; the modulo shortcut would quietly favour the low ones.
func newPairCode() (string, error) {
	const limit = 1000000
	// 4294967295 is not a multiple of 1,000,000, so values in the final
	// partial bucket are rejected and redrawn.
	const bound = 4294000000
	buf := make([]byte, 4)
	for tries := 0; tries < 100; tries++ {
		if _, err := rand.Read(buf); err != nil {
			return "", fmt.Errorf("cannot generate a pairing code: %w", err)
		}
		v := uint32(buf[0])<<24 | uint32(buf[1])<<16 | uint32(buf[2])<<8 | uint32(buf[3])
		if uint64(v) >= bound {
			continue
		}
		return fmt.Sprintf("%06d", v%limit), nil
	}
	return "", errors.New("cannot generate a pairing code")
}

// info returns the device descriptor SnapBeam publishes about itself.
func (rc *receiver) info() deviceInfo {
	return deviceInfo{
		Alias:       rc.Alias,
		Version:     protocolVersion,
		DeviceModel: rc.DeviceModel,
		DeviceType:  rc.DeviceType,
		Fingerprint: rc.Fingerprint,
		Port:        rc.Port,
		Protocol:    "http",
		// SnapBeam is a receiver: the download API (protocol section 5) is
		// not implemented, so this is false and stays false.
		Download: false,
	}
}

func (rc *receiver) registerBody() registerResponse {
	return registerResponse{
		Alias:       rc.Alias,
		Version:     protocolVersion,
		DeviceModel: rc.DeviceModel,
		DeviceType:  rc.DeviceType,
		Fingerprint: rc.Fingerprint,
		Download:    false,
	}
}

// checkPairCode compares a supplied code with the pairing code without leaking
// how much of it was right through the time the comparison took.
//
// subtle.ConstantTimeCompare returns 0 for inputs of different lengths without
// looking at them, so the lengths are equalised first with a SHA-256 of each
// side. That keeps the comparison constant-time for a wrong code of any
// length, which is the case that matters when somebody is guessing.
func checkPairCode(expected, supplied string) bool {
	if expected == "" {
		return true // no code configured; the caller decides whether that is allowed
	}
	e := sha256.Sum256([]byte(expected))
	s := sha256.Sum256([]byte(supplied))
	return subtle.ConstantTimeCompare(e[:], s[:]) == 1
}

// clientIP is the address part of a request's RemoteAddr, used to tie an
// upload to the sender that was granted the session.
func clientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

// checkPIN implements the PIN rules of protocol sections 4.1 and 5.2: 401 when
// the PIN is missing or wrong, 429 once an address has got it wrong three
// times. The code itself is never written to the log.
func (rc *receiver) checkPIN(r *http.Request) (int, string) {
	if rc.PairCode == "" {
		return 0, ""
	}
	ip := clientIP(r)

	rc.mu.Lock()
	fails := rc.pinFails[ip]
	rc.mu.Unlock()
	if fails >= maxPINAttempts {
		return http.StatusTooManyRequests, "too many wrong pairing codes from this address"
	}

	supplied := r.URL.Query().Get("pin")
	if supplied == "" {
		return http.StatusUnauthorized, "pairing code required"
	}
	if !checkPairCode(rc.PairCode, supplied) {
		rc.mu.Lock()
		rc.pinFails[ip]++
		rc.mu.Unlock()
		return http.StatusUnauthorized, "invalid pairing code"
	}

	rc.mu.Lock()
	delete(rc.pinFails, ip)
	rc.mu.Unlock()
	return 0, ""
}

// ---------------------------------------------------------------------------
// Routes
// ---------------------------------------------------------------------------

func writeJSON(w http.ResponseWriter, status int, body any) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(body)
}

// registerLocalSendRoutes attaches the v2 receive API to mux.
func (rc *receiver) registerLocalSendRoutes(mux *http.ServeMux) {
	mux.HandleFunc("POST /api/localsend/v2/register", rc.handleRegister)
	mux.HandleFunc("GET /api/localsend/v2/info", rc.handleInfo)
	mux.HandleFunc("POST /api/localsend/v2/prepare-upload", rc.handlePrepareUpload)
	mux.HandleFunc("POST /api/localsend/v2/upload", rc.handleUpload)
	mux.HandleFunc("POST /api/localsend/v2/cancel", rc.handleCancel)
}

// handleRegister answers the two-way discovery handshake (section 3.1/3.2). It
// is deliberately not PIN-protected: the reference implementation does not
// protect it either, and a device that cannot see us cannot ask us for a PIN.
// It discloses only the alias and fingerprint we are already shouting into the
// multicast group.
func (rc *receiver) handleRegister(w http.ResponseWriter, r *http.Request) {
	var peer deviceInfo
	if err := json.NewDecoder(io.LimitReader(r.Body, 64<<10)).Decode(&peer); err != nil {
		http.Error(w, "invalid body", http.StatusBadRequest)
		return
	}
	if peer.Alias != "" && peer.Fingerprint != rc.Fingerprint {
		if rc.onRegister != nil {
			rc.onRegister(peerFromInfo(peer, clientIP(r)))
		} else {
			fmt.Printf("%s  %q registered from %s\n", stamp(), peer.Alias, clientIP(r))
		}
	}
	writeJSON(w, http.StatusOK, rc.registerBody())
}

func (rc *receiver) handleInfo(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, infoResponse(rc.registerBody()))
}

// handlePrepareUpload implements section 4.1: check the PIN, take the single
// session slot, decide which of the offered files to accept, and hand back one
// unguessable token per accepted file.
func (rc *receiver) handlePrepareUpload(w http.ResponseWriter, r *http.Request) {
	if status, msg := rc.checkPIN(r); status != 0 {
		http.Error(w, msg, status)
		return
	}

	var req prepareUploadRequest
	if err := json.NewDecoder(io.LimitReader(r.Body, maxMetadataBytes)).Decode(&req); err != nil {
		http.Error(w, "invalid body", http.StatusBadRequest)
		return
	}
	if len(req.Files) == 0 {
		http.Error(w, "no files provided", http.StatusBadRequest)
		return
	}

	sessionID, err := newToken()
	if err != nil {
		http.Error(w, "internal error", http.StatusInternalServerError)
		return
	}

	rc.mu.Lock()
	if rc.sess != nil && time.Since(rc.sess.touched) > sessionIdleTimeout {
		rc.sess = nil
	}
	if rc.sess != nil {
		rc.mu.Unlock()
		http.Error(w, "blocked by another session", http.StatusConflict)
		return
	}

	files := map[string]*sessionFile{}
	tokens := map[string]string{}
	var skipped []string
	for id, f := range req.Files {
		if f.Size > rc.MaxBytes {
			skipped = append(skipped, fmt.Sprintf("%s (%s)", safeFileName(f.FileName), humanBytes(f.Size)))
			continue
		}
		tok, err := newToken()
		if err != nil {
			rc.mu.Unlock()
			http.Error(w, "internal error", http.StatusInternalServerError)
			return
		}
		if f.ID == "" {
			f.ID = id
		}
		files[id] = &sessionFile{dto: f, token: tok, status: filePending}
		tokens[id] = tok
	}

	if len(files) == 0 {
		rc.mu.Unlock()
		if len(skipped) > 0 {
			fmt.Printf("%s  refused %d file(s) over the %s limit: %s\n",
				stamp(), len(skipped), humanBytes(rc.MaxBytes), strings.Join(skipped, ", "))
			http.Error(w, "every offered file is over the receiver's size limit of "+humanBytes(rc.MaxBytes),
				http.StatusForbidden)
			return
		}
		// Nothing to do, and nothing wrong: section 4.1 spells this as 204.
		w.WriteHeader(http.StatusNoContent)
		return
	}

	alias := strings.TrimSpace(req.Info.Alias)
	if alias == "" {
		alias = clientIP(r)
	}
	now := time.Now()
	rc.sess = &uploadSession{
		id:       sessionID,
		senderIP: clientIP(r),
		alias:    alias,
		files:    files,
		started:  now,
		touched:  now,
	}
	rc.mu.Unlock()

	if len(skipped) > 0 {
		fmt.Printf("%s  skipped %d file(s) over the %s limit: %s\n",
			stamp(), len(skipped), humanBytes(rc.MaxBytes), strings.Join(skipped, ", "))
	}
	fmt.Printf("%s  incoming from %q via LocalSend: %d file(s)\n", stamp(), alias, len(files))

	writeJSON(w, http.StatusOK, prepareUploadResponse{SessionID: sessionID, Files: tokens})
}

// handleUpload implements section 4.2: validate session, file and token, then
// stream the body to disk. The sender's address must match the one that was
// granted the session, so a token overheard on the network is not enough.
func (rc *receiver) handleUpload(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	sessionID, fileID, token := q.Get("sessionId"), q.Get("fileId"), q.Get("token")
	if sessionID == "" || fileID == "" || token == "" {
		http.Error(w, "missing parameters", http.StatusBadRequest)
		return
	}

	rc.mu.Lock()
	sess := rc.sess
	if sess == nil || sess.id != sessionID || sess.senderIP != clientIP(r) {
		rc.mu.Unlock()
		http.Error(w, "invalid token or IP address", http.StatusForbidden)
		return
	}
	sf, ok := sess.files[fileID]
	if !ok || sf.status != filePending || subtle.ConstantTimeCompare([]byte(sf.token), []byte(token)) != 1 {
		rc.mu.Unlock()
		http.Error(w, "invalid token or IP address", http.StatusForbidden)
		return
	}
	sf.status = fileInProgress
	sess.touched = time.Now()
	dto := sf.dto
	rc.mu.Unlock()

	name := dto.FileName
	if strings.TrimSpace(name) == "" {
		name = "localsend" + extensionFor(dto.FileType)
	}

	// The declared size is a claim, not a fact: the cap is enforced on the
	// bytes that actually arrive.
	hasher := sha256.New()
	item, err := saveStream(rc.Dir, name, io.TeeReader(r.Body, hasher), rc.MaxBytes, time.Now())
	if err != nil {
		rc.mu.Lock()
		sf.status = fileFailed
		rc.finishIfDoneLocked()
		rc.mu.Unlock()
		if errors.Is(err, errTooLarge) {
			fmt.Printf("%s  refused %q: %v\n", stamp(), safeFileName(name), err)
			http.Error(w, err.Error(), http.StatusForbidden)
			return
		}
		fmt.Fprintf(stderr, "%s: %v\n", appName, err)
		http.Error(w, "could not write the file", http.StatusInternalServerError)
		return
	}

	// Section 4.2: if the sender declared a checksum, verify it and answer
	// 422 on a mismatch. The file is kept — nothing is ever deleted — but it
	// is reported as suspect.
	if want := dto.checksum(); want != "" {
		got := hex.EncodeToString(hasher.Sum(nil))
		if !strings.EqualFold(got, want) {
			rc.mu.Lock()
			sf.status = fileFailed
			rc.finishIfDoneLocked()
			rc.mu.Unlock()
			fmt.Printf("%s  CHECKSUM MISMATCH on %s - kept, but the bytes are not what the sender described\n",
				stamp(), item.Path)
			http.Error(w, "checksum mismatch", http.StatusUnprocessableEntity)
			return
		}
	}

	rc.mu.Lock()
	sf.status = fileDone
	sess.touched = time.Now()
	alias := sess.alias
	rc.finishIfDoneLocked()
	rc.mu.Unlock()

	item.Via = "LocalSend"
	item.From = alias
	rc.deliver(item)
	w.WriteHeader(http.StatusOK)
}

// handleCancel implements section 4.3.
func (rc *receiver) handleCancel(w http.ResponseWriter, r *http.Request) {
	sessionID := r.URL.Query().Get("sessionId")
	rc.mu.Lock()
	sess := rc.sess
	if sess != nil && sess.senderIP == clientIP(r) && (sessionID == "" || sessionID == sess.id) {
		rc.sess = nil
		rc.mu.Unlock()
		fmt.Printf("%s  sender cancelled the transfer\n", stamp())
		w.WriteHeader(http.StatusOK)
		return
	}
	rc.mu.Unlock()
	// A cancel for a session we do not manage is not an error worth shouting
	// about: the peer may be cancelling a transfer in the other direction.
	w.WriteHeader(http.StatusOK)
}

// finishIfDoneLocked releases the session slot once every accepted file has
// reached a final state. Callers must hold rc.mu.
func (rc *receiver) finishIfDoneLocked() {
	if rc.sess == nil {
		return
	}
	for _, f := range rc.sess.files {
		if f.status == filePending || f.status == fileInProgress {
			return
		}
	}
	rc.sess = nil
}

// deliver records an item and hands it to the reporter. Both the LocalSend
// path and the browser path end here, which is what makes them print
// identically and land in the same folder.
func (rc *receiver) deliver(item savedItem) {
	rc.mu.Lock()
	rc.received = append(rc.received, item)
	rc.mu.Unlock()
	if rc.report != nil {
		rc.report(item)
	}
}

// sweepSessions releases a session whose sender stopped talking, so one
// crashed phone cannot lock the receiver out with 409 forever.
func (rc *receiver) sweepSessions(stop <-chan struct{}) {
	t := time.NewTicker(30 * time.Second)
	defer t.Stop()
	for {
		select {
		case <-stop:
			return
		case <-t.C:
			rc.mu.Lock()
			if rc.sess != nil && time.Since(rc.sess.touched) > sessionIdleTimeout {
				fmt.Printf("%s  transfer from %q went quiet for %s - session released\n",
					stamp(), rc.sess.alias, sessionIdleTimeout)
				rc.sess = nil
			}
			rc.mu.Unlock()
		}
	}
}