Files
urlshortener/main.go
2026-05-19 06:31:43 +02:00

733 lines
21 KiB
Go

package main
import (
"context"
"crypto/hmac"
"crypto/rand"
"crypto/sha256"
"database/sql"
"embed"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"io/fs"
"log"
"net/http"
"net/url"
"os"
"strings"
"time"
_ "modernc.org/sqlite"
)
//go:embed web/*
var webFiles embed.FS
type AuthMode string
const (
AuthModeLocal AuthMode = "local"
AuthModeOIDC AuthMode = "oidc"
)
type AuthConfig struct {
Mode AuthMode
Issuer string
ClientID string
ClientSecret string
RedirectURL string
Scopes []string
SessionSecret []byte
}
type oidcDiscovery struct {
AuthorizationEndpoint string `json:"authorization_endpoint"`
TokenEndpoint string `json:"token_endpoint"`
UserInfoEndpoint string `json:"userinfo_endpoint"`
Issuer string `json:"issuer"`
}
type oauthTokenResponse struct {
AccessToken string `json:"access_token"`
TokenType string `json:"token_type"`
IDToken string `json:"id_token"`
ExpiresIn int `json:"expires_in"`
}
type userClaims struct {
Subject string `json:"sub"`
PreferredUsername string `json:"preferred_username"`
Name string `json:"name"`
Email string `json:"email"`
}
type sessionData struct {
Username string `json:"username"`
Subject string `json:"sub,omitempty"`
Email string `json:"email,omitempty"`
Expires int64 `json:"exp"`
}
type Server struct {
db *sql.DB
baseURL string
auth AuthConfig
oidc oidcDiscovery
}
type Link struct {
Code string `json:"code"`
ShortURL string `json:"short_url"`
LongURL string `json:"long_url"`
CreatedAt time.Time `json:"created_at"`
}
type CreateRequest struct {
URL string `json:"url"`
}
type UpdateRequest struct {
URL string `json:"url"`
}
type CreateResponse struct {
Code string `json:"code"`
ShortURL string `json:"short_url"`
LongURL string `json:"long_url"`
}
type UserResponse struct {
Authenticated bool `json:"authenticated"`
Username string `json:"username,omitempty"`
Email string `json:"email,omitempty"`
AuthMode string `json:"auth_mode"`
}
type ErrorResponse struct {
Error string `json:"error"`
}
func main() {
addr := env("ADDR", ":8081")
dbPath := env("DB_PATH", "/data/shortener.db")
baseURL := strings.TrimRight(env("BASE_URL", "http://localhost:8081"), "/")
auth, err := authConfigFromEnv()
if err != nil {
log.Fatal(err)
}
db, err := sql.Open("sqlite", dbPath)
if err != nil {
log.Fatal(err)
}
defer db.Close()
if err := migrate(db); err != nil {
log.Fatal(err)
}
s := &Server{db: db, baseURL: baseURL, auth: auth}
if err := s.discoverOIDC(); err != nil {
log.Fatal(err)
}
mux := http.NewServeMux()
mux.HandleFunc("GET /healthz", s.health)
mux.HandleFunc("GET /api/me", s.me)
mux.HandleFunc("GET /api/links", s.requireAuth(s.listLinks))
mux.HandleFunc("POST /api/shorten", s.requireAuth(s.createShortURL))
mux.HandleFunc("PUT /api/links/{code}", s.requireAuth(s.updateLink))
mux.HandleFunc("DELETE /api/links/{code}", s.requireAuth(s.deleteLink))
mux.HandleFunc("GET /auth/login", s.handleLogin)
mux.HandleFunc("GET /auth/callback", s.handleOIDCCallback)
mux.HandleFunc("POST /auth/logout", s.handleLogout)
staticFiles, err := fs.Sub(webFiles, "web")
if err != nil {
log.Fatal(err)
}
mux.Handle("GET /assets/", http.FileServerFS(staticFiles))
mux.HandleFunc("GET /", s.homeOrRedirect(staticFiles))
log.Printf("URL shortener listening on %s", addr)
log.Printf("Base URL: %s", baseURL)
log.Printf("Auth mode: %s", auth.Mode)
if err := http.ListenAndServe(addr, logging(mux)); err != nil {
log.Fatal(err)
}
}
func migrate(db *sql.DB) error {
_, err := db.Exec(`
CREATE TABLE IF NOT EXISTS links (
code TEXT PRIMARY KEY,
long_url TEXT NOT NULL,
owner_sub TEXT NOT NULL DEFAULT '',
owner_name TEXT NOT NULL DEFAULT '',
created_at DATETIME NOT NULL,
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS idx_links_owner_sub ON links(owner_sub);
CREATE INDEX IF NOT EXISTS idx_links_created_at ON links(created_at);
`)
if err != nil {
return err
}
_ = addColumnIfMissing(db, "links", "owner_sub", "TEXT NOT NULL DEFAULT ''")
_ = addColumnIfMissing(db, "links", "owner_name", "TEXT NOT NULL DEFAULT ''")
_ = addColumnIfMissing(db, "links", "updated_at", "DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP")
_, _ = db.Exec(`CREATE INDEX IF NOT EXISTS idx_links_owner_sub ON links(owner_sub)`)
return nil
}
func addColumnIfMissing(db *sql.DB, table, column, definition string) error {
rows, err := db.Query(`PRAGMA table_info(` + table + `)`)
if err != nil {
return err
}
defer rows.Close()
for rows.Next() {
var cid int
var name, typ string
var notNull int
var dflt any
var pk int
if err := rows.Scan(&cid, &name, &typ, &notNull, &dflt, &pk); err != nil {
return err
}
if name == column {
return nil
}
}
_, err = db.Exec(`ALTER TABLE ` + table + ` ADD COLUMN ` + column + ` ` + definition)
return err
}
func (s *Server) health(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}
func (s *Server) me(w http.ResponseWriter, r *http.Request) {
session, ok := s.readSession(r)
if !ok {
writeJSON(w, http.StatusOK, UserResponse{Authenticated: false, AuthMode: string(s.auth.Mode)})
return
}
writeJSON(w, http.StatusOK, UserResponse{Authenticated: true, Username: session.Username, Email: session.Email, AuthMode: string(s.auth.Mode)})
}
func (s *Server) listLinks(w http.ResponseWriter, r *http.Request, session sessionData) {
rows, err := s.db.Query(`
SELECT code, long_url, created_at
FROM links
WHERE owner_sub = ?
ORDER BY created_at DESC
LIMIT 200`, session.Subject)
if err != nil {
log.Printf("list links failed: %v", err)
writeError(w, http.StatusInternalServerError, "could not load links")
return
}
defer rows.Close()
links := make([]Link, 0)
for rows.Next() {
var l Link
if err := rows.Scan(&l.Code, &l.LongURL, &l.CreatedAt); err != nil {
writeError(w, http.StatusInternalServerError, "could not read links")
return
}
l.ShortURL = s.baseURL + "/" + l.Code
links = append(links, l)
}
writeJSON(w, http.StatusOK, links)
}
func (s *Server) createShortURL(w http.ResponseWriter, r *http.Request, session sessionData) {
var req CreateRequest
if err := json.NewDecoder(io.LimitReader(r.Body, 1<<20)).Decode(&req); err != nil {
writeError(w, http.StatusBadRequest, "invalid JSON body")
return
}
longURL := strings.TrimSpace(req.URL)
if !validHTTPURL(longURL) {
writeError(w, http.StatusBadRequest, "url must be a valid http or https URL")
return
}
code, err := s.insertWithRandomCode(longURL, session)
if err != nil {
log.Printf("create short url failed: %v", err)
writeError(w, http.StatusInternalServerError, "could not create short URL")
return
}
writeJSON(w, http.StatusCreated, CreateResponse{Code: code, ShortURL: s.baseURL + "/" + code, LongURL: longURL})
}
func (s *Server) updateLink(w http.ResponseWriter, r *http.Request, session sessionData) {
code := strings.TrimSpace(r.PathValue("code"))
var req UpdateRequest
if err := json.NewDecoder(io.LimitReader(r.Body, 1<<20)).Decode(&req); err != nil {
writeError(w, http.StatusBadRequest, "invalid JSON body")
return
}
longURL := strings.TrimSpace(req.URL)
if !validHTTPURL(longURL) {
writeError(w, http.StatusBadRequest, "url must be a valid http or https URL")
return
}
res, err := s.db.Exec(`
UPDATE links
SET long_url = ?, updated_at = ?
WHERE code = ? AND owner_sub = ?`, longURL, time.Now().UTC(), code, session.Subject)
if err != nil {
log.Printf("update link failed: %v", err)
writeError(w, http.StatusInternalServerError, "could not update link")
return
}
affected, _ := res.RowsAffected()
if affected == 0 {
writeError(w, http.StatusNotFound, "link not found")
return
}
writeJSON(w, http.StatusOK, CreateResponse{Code: code, ShortURL: s.baseURL + "/" + code, LongURL: longURL})
}
func (s *Server) deleteLink(w http.ResponseWriter, r *http.Request, session sessionData) {
code := strings.TrimSpace(r.PathValue("code"))
res, err := s.db.Exec(`DELETE FROM links WHERE code = ? AND owner_sub = ?`, code, session.Subject)
if err != nil {
log.Printf("delete link failed: %v", err)
writeError(w, http.StatusInternalServerError, "could not delete link")
return
}
affected, _ := res.RowsAffected()
if affected == 0 {
writeError(w, http.StatusNotFound, "link not found")
return
}
w.WriteHeader(http.StatusNoContent)
}
func (s *Server) homeOrRedirect(staticFiles fs.FS) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
path := strings.Trim(r.URL.Path, "/")
if path == "" {
serveIndex(w, r, staticFiles)
return
}
if path == "favicon.ico" || strings.HasPrefix(path, "assets/") || strings.HasPrefix(path, "api/") || strings.HasPrefix(path, "auth/") {
http.NotFound(w, r)
return
}
s.redirect(w, r, path)
}
}
func serveIndex(w http.ResponseWriter, r *http.Request, staticFiles fs.FS) {
w.Header().Set("Content-Type", "text/html; charset=utf-8")
http.ServeFileFS(w, r, staticFiles, "index.html")
}
func (s *Server) redirect(w http.ResponseWriter, r *http.Request, code string) {
var longURL string
err := s.db.QueryRow(`SELECT long_url FROM links WHERE code = ?`, code).Scan(&longURL)
if errors.Is(err, sql.ErrNoRows) {
http.NotFound(w, r)
return
}
if err != nil {
log.Printf("lookup failed: %v", err)
writeError(w, http.StatusInternalServerError, "lookup failed")
return
}
http.Redirect(w, r, longURL, http.StatusFound)
}
func (s *Server) insertWithRandomCode(longURL string, session sessionData) (string, error) {
for range 8 {
code, err := randomCode(7)
if err != nil {
return "", err
}
now := time.Now().UTC()
_, err = s.db.Exec(`INSERT INTO links(code, long_url, owner_sub, owner_name, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)`, code, longURL, session.Subject, session.Username, now, now)
if err == nil {
return code, nil
}
if !strings.Contains(strings.ToLower(err.Error()), "constraint") {
return "", err
}
}
return "", errors.New("could not generate unique code")
}
func (s *Server) requireAuth(next func(http.ResponseWriter, *http.Request, sessionData)) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
session, ok := s.readSession(r)
if !ok {
writeError(w, http.StatusUnauthorized, "login required")
return
}
next(w, r, session)
}
}
func authConfigFromEnv() (AuthConfig, error) {
mode := AuthMode(strings.ToLower(strings.TrimSpace(env("AUTH_MODE", "local"))))
if mode != AuthModeLocal && mode != AuthModeOIDC {
return AuthConfig{}, fmt.Errorf("unsupported AUTH_MODE %q", mode)
}
secret := strings.TrimSpace(os.Getenv("SESSION_SECRET"))
var secretBytes []byte
if secret == "" {
secretBytes = make([]byte, 32)
if _, err := rand.Read(secretBytes); err != nil {
return AuthConfig{}, err
}
log.Printf("SESSION_SECRET is not set; generated a temporary secret. Existing sessions will be invalid after restart")
} else {
decoded, err := base64.StdEncoding.DecodeString(secret)
if err == nil && len(decoded) >= 32 {
secretBytes = decoded
} else {
secretBytes = []byte(secret)
}
if len(secretBytes) < 32 {
return AuthConfig{}, errors.New("SESSION_SECRET must be at least 32 bytes or a base64 encoded 32 byte value")
}
}
cfg := AuthConfig{
Mode: mode,
Issuer: strings.TrimRight(strings.TrimSpace(os.Getenv("OIDC_ISSUER")), "/"),
ClientID: strings.TrimSpace(os.Getenv("OIDC_CLIENT_ID")),
ClientSecret: strings.TrimSpace(os.Getenv("OIDC_CLIENT_SECRET")),
RedirectURL: strings.TrimSpace(os.Getenv("OIDC_REDIRECT_URL")),
Scopes: splitScopes(env("OIDC_SCOPES", "openid profile email")),
SessionSecret: secretBytes,
}
if cfg.Mode == AuthModeOIDC {
if cfg.Issuer == "" || cfg.ClientID == "" || cfg.ClientSecret == "" || cfg.RedirectURL == "" {
return AuthConfig{}, errors.New("AUTH_MODE=oidc requires OIDC_ISSUER, OIDC_CLIENT_ID, OIDC_CLIENT_SECRET and OIDC_REDIRECT_URL")
}
}
return cfg, nil
}
func (s *Server) discoverOIDC() error {
if s.auth.Mode != AuthModeOIDC {
return nil
}
discoveryURL := s.auth.Issuer + "/.well-known/openid-configuration"
req, err := http.NewRequest(http.MethodGet, discoveryURL, nil)
if err != nil {
return err
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
return fmt.Errorf("OIDC discovery failed: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode > 299 {
return fmt.Errorf("OIDC discovery returned %s", resp.Status)
}
var d oidcDiscovery
if err := json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&d); err != nil {
return err
}
if d.AuthorizationEndpoint == "" || d.TokenEndpoint == "" || d.UserInfoEndpoint == "" {
return errors.New("OIDC discovery response is missing required endpoints")
}
s.oidc = d
return nil
}
func (s *Server) handleLogin(w http.ResponseWriter, r *http.Request) {
if s.auth.Mode == AuthModeLocal {
_ = s.setSession(w, r, sessionData{Username: "local-user", Subject: "local-user", Expires: time.Now().Add(30 * 24 * time.Hour).Unix()})
http.Redirect(w, r, "/", http.StatusSeeOther)
return
}
s.handleOIDCLogin(w, r)
}
func (s *Server) handleOIDCLogin(w http.ResponseWriter, r *http.Request) {
state, err := randomString(32)
if err != nil {
http.Error(w, "could not create login state", http.StatusInternalServerError)
return
}
setCookie(w, &http.Cookie{Name: "oauth_state", Value: state, MaxAge: 10 * 60, Path: "/auth", HttpOnly: true, SameSite: http.SameSiteLaxMode, Secure: isSecure(r)})
q := url.Values{}
q.Set("response_type", "code")
q.Set("client_id", s.auth.ClientID)
q.Set("redirect_uri", s.auth.RedirectURL)
q.Set("scope", strings.Join(s.auth.Scopes, " "))
q.Set("state", state)
http.Redirect(w, r, s.oidc.AuthorizationEndpoint+"?"+q.Encode(), http.StatusFound)
}
func (s *Server) handleOIDCCallback(w http.ResponseWriter, r *http.Request) {
if s.auth.Mode != AuthModeOIDC {
http.Redirect(w, r, "/", http.StatusSeeOther)
return
}
if errText := r.URL.Query().Get("error"); errText != "" {
http.Error(w, "login cancelled: "+errText, http.StatusUnauthorized)
return
}
stateCookie, err := r.Cookie("oauth_state")
if err != nil || stateCookie.Value == "" || stateCookie.Value != r.URL.Query().Get("state") {
http.Error(w, "invalid oauth state", http.StatusBadRequest)
return
}
clearCookie(w, "oauth_state", "/auth")
code := r.URL.Query().Get("code")
if code == "" {
http.Error(w, "missing authorization code", http.StatusBadRequest)
return
}
tok, err := s.exchangeCode(r.Context(), code)
if err != nil {
log.Printf("oauth token exchange failed: %v", err)
http.Error(w, "token exchange failed", http.StatusBadGateway)
return
}
claims, err := s.fetchUserInfo(r.Context(), tok.AccessToken)
if err != nil {
log.Printf("oauth userinfo failed: %v", err)
http.Error(w, "could not fetch user info", http.StatusBadGateway)
return
}
username := normalizeUsername(claims)
if username == "" || strings.TrimSpace(claims.Subject) == "" {
http.Error(w, "identity provider did not return usable identity data", http.StatusBadGateway)
return
}
if err := s.setSession(w, r, sessionData{Username: username, Subject: claims.Subject, Email: claims.Email, Expires: time.Now().Add(30 * 24 * time.Hour).Unix()}); err != nil {
http.Error(w, "could not create session", http.StatusInternalServerError)
return
}
http.Redirect(w, r, "/", http.StatusSeeOther)
}
func (s *Server) exchangeCode(ctx context.Context, code string) (oauthTokenResponse, error) {
values := url.Values{}
values.Set("grant_type", "authorization_code")
values.Set("code", code)
values.Set("redirect_uri", s.auth.RedirectURL)
values.Set("client_id", s.auth.ClientID)
values.Set("client_secret", s.auth.ClientSecret)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.oidc.TokenEndpoint, strings.NewReader(values.Encode()))
if err != nil {
return oauthTokenResponse{}, err
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
req.Header.Set("Accept", "application/json")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return oauthTokenResponse{}, err
}
defer resp.Body.Close()
body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
if err != nil {
return oauthTokenResponse{}, err
}
if resp.StatusCode < 200 || resp.StatusCode > 299 {
return oauthTokenResponse{}, fmt.Errorf("token endpoint returned %s: %s", resp.Status, string(body))
}
var tok oauthTokenResponse
if err := json.Unmarshal(body, &tok); err != nil {
return oauthTokenResponse{}, err
}
if tok.AccessToken == "" {
return oauthTokenResponse{}, errors.New("token endpoint did not return access_token")
}
return tok, nil
}
func (s *Server) fetchUserInfo(ctx context.Context, accessToken string) (userClaims, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.oidc.UserInfoEndpoint, nil)
if err != nil {
return userClaims{}, err
}
req.Header.Set("Authorization", "Bearer "+accessToken)
req.Header.Set("Accept", "application/json")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return userClaims{}, err
}
defer resp.Body.Close()
body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
if err != nil {
return userClaims{}, err
}
if resp.StatusCode < 200 || resp.StatusCode > 299 {
return userClaims{}, fmt.Errorf("userinfo endpoint returned %s: %s", resp.Status, string(body))
}
var claims userClaims
if err := json.Unmarshal(body, &claims); err != nil {
return userClaims{}, err
}
return claims, nil
}
func (s *Server) handleLogout(w http.ResponseWriter, r *http.Request) {
clearCookie(w, "shortener_session", "/")
w.WriteHeader(http.StatusNoContent)
}
func (s *Server) setSession(w http.ResponseWriter, r *http.Request, data sessionData) error {
b, err := json.Marshal(data)
if err != nil {
return err
}
payload := base64.RawURLEncoding.EncodeToString(b)
sig := sign(payload, s.auth.SessionSecret)
setCookie(w, &http.Cookie{Name: "shortener_session", Value: payload + "." + sig, Path: "/", MaxAge: 60 * 60 * 24 * 30, HttpOnly: true, SameSite: http.SameSiteLaxMode, Secure: isSecure(r)})
return nil
}
func (s *Server) readSession(r *http.Request) (sessionData, bool) {
c, err := r.Cookie("shortener_session")
if err != nil || c.Value == "" {
return sessionData{}, false
}
parts := strings.Split(c.Value, ".")
if len(parts) != 2 || !verify(parts[0], parts[1], s.auth.SessionSecret) {
return sessionData{}, false
}
b, err := base64.RawURLEncoding.DecodeString(parts[0])
if err != nil {
return sessionData{}, false
}
var data sessionData
if err := json.Unmarshal(b, &data); err != nil {
return sessionData{}, false
}
if data.Expires < time.Now().Unix() || strings.TrimSpace(data.Username) == "" || strings.TrimSpace(data.Subject) == "" {
return sessionData{}, false
}
return data, true
}
func randomCode(length int) (string, error) {
b := make([]byte, length)
if _, err := rand.Read(b); err != nil {
return "", err
}
code := base64.RawURLEncoding.EncodeToString(b)
if len(code) > length {
code = code[:length]
}
return code, nil
}
func validHTTPURL(raw string) bool {
u, err := url.ParseRequestURI(raw)
if err != nil {
return false
}
return (u.Scheme == "http" || u.Scheme == "https") && u.Host != ""
}
func writeJSON(w http.ResponseWriter, status int, v any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(v)
}
func writeError(w http.ResponseWriter, status int, msg string) {
writeJSON(w, status, ErrorResponse{Error: msg})
}
func env(key, fallback string) string {
if value := os.Getenv(key); value != "" {
return value
}
return fallback
}
func splitScopes(s string) []string {
fields := strings.Fields(s)
if len(fields) == 0 {
return []string{"openid", "profile", "email"}
}
return fields
}
func sign(payload string, secret []byte) string {
mac := hmac.New(sha256.New, secret)
_, _ = mac.Write([]byte(payload))
return base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}
func verify(payload, got string, secret []byte) bool {
expected := sign(payload, secret)
return hmac.Equal([]byte(expected), []byte(got))
}
func randomString(n int) (string, error) {
b := make([]byte, n)
if _, err := rand.Read(b); err != nil {
return "", err
}
return base64.RawURLEncoding.EncodeToString(b), nil
}
func normalizeUsername(c userClaims) string {
for _, v := range []string{c.PreferredUsername, c.Name, c.Email, c.Subject} {
v = strings.TrimSpace(v)
if v != "" {
if strings.Contains(v, "@") {
v = strings.Split(v, "@")[0]
}
if len(v) > 40 {
v = v[:40]
}
return v
}
}
return ""
}
func setCookie(w http.ResponseWriter, c *http.Cookie) { http.SetCookie(w, c) }
func clearCookie(w http.ResponseWriter, name, path string) {
http.SetCookie(w, &http.Cookie{Name: name, Value: "", Path: path, MaxAge: -1, HttpOnly: true, SameSite: http.SameSiteLaxMode})
}
func isSecure(r *http.Request) bool {
return r.TLS != nil || strings.EqualFold(r.Header.Get("X-Forwarded-Proto"), "https")
}
func logging(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
start := time.Now()
next.ServeHTTP(w, r)
log.Printf("%s %s %s", r.Method, r.URL.Path, time.Since(start).Round(time.Millisecond))
})
}