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, ¬Null, &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)) }) }