handler.go

package main

import (
	"context"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/acoshift/pgsql/pgctx"
	"gocloud.dev/blob"
	"gocloud.dev/gcerrors"
)

type App struct {
	Bucket         *blob.Bucket
	BaseURL        string
	CDNBaseURL     string
	InternalSecret string
	SignKey        []byte
	// MaxUploadSize caps the maxSize a signed-upload URL will accept (bytes).
	// Zero falls back to defaultMaxUploadSize.
	MaxUploadSize int64
	checkAuth     func(ctx context.Context, auth, project, projectID string) AuthResult
}

func (a *App) routes() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /{$}", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Deploys.app Dropbox Service"))
	})
	mux.HandleFunc("POST /{$}", a.uploadHandler)
	mux.HandleFunc("POST /uploads", a.uploadURLHandler)
	mux.HandleFunc("PUT /uploads/{token}", a.uploadDirectHandler)
	mux.HandleFunc("GET /files/{token}", a.fileHandler)
	mux.HandleFunc("GET /_cdn/{token}", a.cdnFileHandler)
	mux.HandleFunc("POST /internal/gc", a.gcHandler)
	return mux
}

func (a *App) auth(ctx context.Context, auth, project, projectID string) AuthResult {
	if a.checkAuth != nil {
		return a.checkAuth(ctx, auth, project, projectID)
	}
	return checkAuth(ctx, auth, project, projectID)
}

func (a *App) uploadHandler(w http.ResponseWriter, r *http.Request) {
	auth := r.Header.Get("Authorization")
	project := firstNonEmpty(r.URL.Query().Get("project"), r.Header.Get("param-project"))
	projectID := firstNonEmpty(r.URL.Query().Get("projectId"), r.Header.Get("param-project-id"))

	// `project` accepts either a sid or a numeric project ID. A project sid
	// always starts with a letter (api ReValidSID: `^[a-z][a-z0-9\-]*[^\-]$`),
	// so an all-digit value is unambiguously a numeric ID — route it to
	// projectID, which me.authorized resolves by ID (a sid lookup would fail).
	// An explicit ?projectId= still wins.
	if projectID == "" && isAllDigits(project) {
		projectID, project = project, ""
	}

	authResult := a.auth(r.Context(), auth, project, projectID)
	if !authResult.Authorized {
		jsonFail(w, "api: unauthorized", http.StatusOK)
		return
	}

	ttlStr := firstNonEmpty(r.URL.Query().Get("ttl"), r.Header.Get("param-ttl"))
	ttlDays, _ := strconv.Atoi(ttlStr)
	if ttlDays < 1 || ttlDays > 365 {
		ttlDays = 1
	}

	filename := firstNonEmpty(r.URL.Query().Get("filename"), r.Header.Get("param-filename"))

	if r.Body == nil || r.ContentLength == 0 {
		jsonFail(w, "body empty", http.StatusOK)
		return
	}

	expiresAt := time.Now().UTC().Add(time.Duration(ttlDays) * 24 * time.Hour)
	fn := generateFilename()
	// The signed token is the public URL component. Persist it so the api's
	// dropbox.List can rebuild download URLs without holding SignKey.
	token := makeToken(a.SignKey, fn)

	opts := &blob.WriterOptions{
		CacheControl: "public, max-age=86400",
	}
	if filename != "" {
		opts.ContentDisposition = fmt.Sprintf(`attachment; filename="%s"`, escapeFilename(filename))
	}

	bw, err := a.Bucket.NewWriter(r.Context(), fn, opts)
	if err != nil {
		slog.Error("upload file", "error", err)
		jsonFail(w, "failed to upload", http.StatusInternalServerError)
		return
	}
	n, err := io.Copy(bw, r.Body)
	if err != nil {
		_ = bw.Close()
		// Close finalizes whatever bytes already streamed, so a mid-stream
		// failure can leave a partial (or 0-byte) object behind with no DB row.
		// Remove it for the same reason as the empty-upload case below — GC is
		// DB-driven and would never reclaim a rowless object. Delete on a
		// cancellation-detached context: a copy error is usually a client
		// disconnect, which already canceled r.Context(), and reusing it here
		// would fail the cleanup and leak the partial object.
		delCtx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), 10*time.Second)
		if derr := a.Bucket.Delete(delCtx, fn); derr != nil && gcerrors.Code(derr) != gcerrors.NotFound {
			slog.Error("delete partial upload", "fn", fn, "error", derr)
		}
		cancel()
		slog.Error("upload file", "error", err)
		jsonFail(w, "failed to upload", http.StatusInternalServerError)
		return
	}
	if err := bw.Close(); err != nil {
		slog.Error("finalize upload", "error", err)
		jsonFail(w, "failed to upload", http.StatusInternalServerError)
		return
	}

	// Reject empty uploads. The r.ContentLength == 0 guard above rejects the
	// common case up front, but a request with an unknown length (chunked
	// transfer encoding, ContentLength == -1) can still carry an empty body —
	// that only becomes apparent after io.Copy reports n == 0. Storing a
	// 0-byte object serves no one and would linger in the bucket, so delete
	// the object we just finalized and bail before writing any DB row (GC is
	// DB-driven and would never reclaim a bucket object with no row). Use
	// deleteObject, which detaches from r.Context(): a client that disconnects
	// right after the empty body cancels r.Context(), and reusing it here would
	// fail the cleanup and strand the rowless object (same reason as the
	// mid-stream error path above).
	if n == 0 {
		a.deleteObject(r, fn)
		jsonFail(w, "body empty", http.StatusOK)
		return
	}

	uploadCount.WithLabelValues(authResult.Project.ID).Inc()
	uploadBytes.WithLabelValues(authResult.Project.ID).Add(float64(n))

	if _, err := pgctx.Exec(r.Context(), `
		INSERT INTO files (fn, project_id, size, filename, ttl, expires_at, token)
		VALUES ($1, $2, $3, $4, $5, $6, $7)
	`, fn, authResult.Project.ID, n, filename, ttlDays, expiresAt, token); err != nil {
		slog.Error("insert file metadata", "error", err)
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]any{
		"ok": true,
		"result": map[string]any{
			"downloadUrl": a.BaseURL + token,
			"expiresAt":   expiresAt.Format(time.RFC3339),
		},
	})
}

// Filename + signature scheme.
//
// `fn` is what we store in the bucket and the DB — a random 24-char
// alphanumeric string (~143 bits of entropy, drawn from
// [0-9A-Za-z] with rejection sampling). No special chars, so URLs are
// clean and we don't need URL-safe-base64 tricks.
//
// `token` is what appears in the public URL: `fn` + "-" + 20-char
// lowercase-hex HMAC-SHA256 tag (80-bit) keyed by App.SignKey. The
// handlers check the tag before any DB or GCS work, so a flood of
// random `/files/{token}` requests by an attacker who doesn't know
// the key gets 404'd in CPU. The "-" separator means we can change
// `fnLen` later without invalidating old tokens — parseToken splits
// on it instead of relying on a fixed cut point. fn is alphanumeric
// and sig is hex, so neither side can contain a "-".
const (
	fnLen    = 24
	sigLen   = 20 // 80 bits of HMAC, hex-encoded
	tokenSep = "-"
)

const fnAlphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"

// generateFilename returns a fresh fnLen-char random alphanumeric fn.
// Uses rejection sampling so each char is uniformly drawn from
// fnAlphabet (avoiding the mod-bias you'd get from `b%62`).
func generateFilename() string {
	out := make([]byte, fnLen)
	buf := make([]byte, fnLen*2) // slack for rejection
	pos := 0
	for pos < fnLen {
		if _, err := rand.Read(buf); err != nil {
			panic(err) // crypto/rand failing is unrecoverable
		}
		for _, b := range buf {
			if pos >= fnLen {
				break
			}
			// 62 * 4 = 248 is the largest unbiased range under 256;
			// reject 248..255.
			if b < 248 {
				out[pos] = fnAlphabet[b%62]
				pos++
			}
		}
	}
	return string(out)
}

// signFilename returns the sigLen-char lowercase-hex HMAC tag for fn
// under key. Caller's responsibility to keep `key` secret.
func signFilename(key []byte, fn string) string {
	mac := hmac.New(sha256.New, key)
	mac.Write([]byte(fn))
	sum := mac.Sum(nil)
	return hex.EncodeToString(sum[:sigLen/2])
}

// makeToken builds the public URL token for fn: `fn + "-" + sig`.
func makeToken(key []byte, fn string) string {
	return fn + tokenSep + signFilename(key, fn)
}

// parseToken splits a public URL token on tokenSep, verifies the HMAC
// in constant time, and returns the fn on success. On any failure —
// missing separator, empty side, bad sig, anything — it returns
// (_, false) and the handler 404s without touching the DB or GCS.
// This is the primary DDoS shield against random `/files/{garbage}`
// floods. The split is structural (on the separator) rather than
// positional, so changing fnLen later won't invalidate old tokens
// still in circulation.
func parseToken(key []byte, token string) (string, bool) {
	idx := strings.IndexByte(token, tokenSep[0])
	if idx <= 0 || idx == len(token)-1 {
		return "", false
	}
	fn := token[:idx]
	sig := token[idx+1:]
	expected := signFilename(key, fn)
	if !hmac.Equal([]byte(sig), []byte(expected)) {
		return "", false
	}
	return fn, true
}

func escapeFilename(s string) string {
	return strings.ReplaceAll(s, `"`, "")
}

func firstNonEmpty(vals ...string) string {
	for _, v := range vals {
		if v != "" {
			return v
		}
	}
	return ""
}

// isAllDigits reports whether s is non-empty and made up solely of ASCII
// digits. It distinguishes a numeric project ID from a project sid: sids
// always start with a letter, so they can never be all-digits.
func isAllDigits(s string) bool {
	if s == "" {
		return false
	}
	for i := 0; i < len(s); i++ {
		if s[i] < '0' || s[i] > '9' {
			return false
		}
	}
	return true
}

// isInternalClient returns true when the request originates from a
// private/loopback/link-local IP per the X-Real-Ip header. Internal callers
// (in-cluster fetches) bypass the CDN redirect and read files directly.
func isInternalClient(r *http.Request) bool {
	ip := net.ParseIP(r.Header.Get("X-Real-Ip"))
	if ip == nil {
		return false
	}
	return ip.IsPrivate() || ip.IsLoopback() || ip.IsLinkLocalUnicast()
}

func jsonFail(w http.ResponseWriter, msg string, status int) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(map[string]any{
		"ok": false,
		"error": map[string]any{
			"message": msg,
		},
	})
}