Project: URL Shortener API

shortener/
    go.mod
    store.go          the data: links in a map, protected by a mutex
    handlers.go       HTTP handlers, JSON helpers, validation, middleware
    main.go           server setup and graceful shutdown
    app_test.go       tests

mkdir shortener
cd shortener
go mod init example.com/shortener

// store.go
package main

import (
	"cmp"
	"slices"
	"sync"
)

const alphabet = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

// encodeBase62 writes n using the 62 characters of alphabet.
func encodeBase62(n int) string {
	if n == 0 {
		return "0"
	}
	var buf []byte
	for n > 0 {
		buf = append(buf, alphabet[n%62])
		n /= 62
	}
	slices.Reverse(buf)
	return string(buf)
}

type Link struct {
	ID     int
	Code   string
	URL    string
	Visits int
}

type Store struct {
	mu     sync.Mutex
	nextID int
	links  map[string]Link
}

func NewStore() *Store {
	return &Store{nextID: 100000, links: make(map[string]Link)}
}

func (s *Store) Add(rawURL string) Link {
	s.mu.Lock()
	defer s.mu.Unlock()
	l := Link{ID: s.nextID, Code: encodeBase62(s.nextID), URL: rawURL}
	s.nextID++
	s.links[l.Code] = l
	return l
}

func (s *Store) Get(code string) (Link, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	l, ok := s.links[code]
	return l, ok
}

// Visit records one visit and returns the long URL.
func (s *Store) Visit(code string) (string, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	l, ok := s.links[code]
	if !ok {
		return "", false
	}
	l.Visits++
	s.links[code] = l
	return l.URL, true
}

func (s *Store) Delete(code string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	_, ok := s.links[code]
	delete(s.links, code)
	return ok
}

// List returns all links ordered by creation.
func (s *Store) List() []Link {
	s.mu.Lock()
	out := make([]Link, 0, len(s.links))
	for _, l := range s.links {
		out = append(out, l)
	}
	s.mu.Unlock()
	slices.SortFunc(out, func(a, b Link) int { return cmp.Compare(a.ID, b.ID) })
	return out
}

// handlers.go (part 1)
package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"time"
)

func validateURL(raw string) error {
	if len(raw) > 2048 {
		return errors.New("url is too long")
	}
	u, err := url.Parse(raw)
	if err != nil {
		return errors.New("url could not be parsed")
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		return errors.New("url must start with http:// or https://")
	}
	if u.Host == "" {
		return errors.New("url must include a host")
	}
	return nil
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	if err := json.NewEncoder(w).Encode(v); err != nil {
		log.Println("write response:", err)
	}
}

func writeError(w http.ResponseWriter, status int, msg string) {
	writeJSON(w, status, map[string]string{"error": msg})
}

func logRequests(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 %v", r.Method, r.URL.Path, time.Since(start))
	})
}

// handlers.go (part 2)
type App struct {
	store   *Store
	baseURL string
}

type linkJSON struct {
	Code     string `json:"code"`
	URL      string `json:"url"`
	ShortURL string `json:"short_url"`
	Visits   int    `json:"visits"`
}

func (a *App) toJSON(l Link) linkJSON {
	return linkJSON{Code: l.Code, URL: l.URL, ShortURL: a.baseURL + "/" + l.Code, Visits: l.Visits}
}

func (a *App) routes() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /{$}", a.home)
	mux.HandleFunc("GET /{code}", a.redirect)
	mux.HandleFunc("POST /api/links", a.createLink)
	mux.HandleFunc("GET /api/links", a.listLinks)
	mux.HandleFunc("GET /api/links/{code}", a.getLink)
	mux.HandleFunc("DELETE /api/links/{code}", a.deleteLink)
	return mux
}

func (a *App) home(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintln(w, "URL shortener is running. POST /api/links to create a link.")
}

func (a *App) createLink(w http.ResponseWriter, r *http.Request) {
	var in struct {
		URL string `json:"url"`
	}
	dec := json.NewDecoder(http.MaxBytesReader(w, r.Body, 4096))
	dec.DisallowUnknownFields()
	if err := dec.Decode(&in); err != nil {
		writeError(w, http.StatusBadRequest, "request body must be JSON with a url field")
		return
	}
	if err := validateURL(in.URL); err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	link := a.store.Add(in.URL)
	w.Header().Set("Location", "/api/links/"+link.Code)
	writeJSON(w, http.StatusCreated, a.toJSON(link))
}

func (a *App) listLinks(w http.ResponseWriter, r *http.Request) {
	links := a.store.List()
	out := make([]linkJSON, 0, len(links))
	for _, l := range links {
		out = append(out, a.toJSON(l))
	}
	writeJSON(w, http.StatusOK, out)
}

func (a *App) getLink(w http.ResponseWriter, r *http.Request) {
	l, ok := a.store.Get(r.PathValue("code"))
	if !ok {
		writeError(w, http.StatusNotFound, "link not found")
		return
	}
	writeJSON(w, http.StatusOK, a.toJSON(l))
}

func (a *App) deleteLink(w http.ResponseWriter, r *http.Request) {
	if !a.store.Delete(r.PathValue("code")) {
		writeError(w, http.StatusNotFound, "link not found")
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (a *App) redirect(w http.ResponseWriter, r *http.Request) {
	target, ok := a.store.Visit(r.PathValue("code"))
	if !ok {
		http.NotFound(w, r)
		return
	}
	http.Redirect(w, r, target, http.StatusFound)
}

// main.go
package main

import (
	"context"
	"errors"
	"log"
	"net/http"
	"os"
	"os/signal"
	"time"
)

func main() {
	addr := ":8080"
	if port, ok := os.LookupEnv("PORT"); ok {
		addr = ":" + port
	}

	app := &App{store: NewStore(), baseURL: "http://localhost" + addr}
	srv := &http.Server{
		Addr:              addr,
		Handler:           logRequests(app.routes()),
		ReadHeaderTimeout: 5 * time.Second,
		ReadTimeout:       10 * time.Second,
		WriteTimeout:      10 * time.Second,
		IdleTimeout:       60 * time.Second,
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()

	go func() {
		log.Printf("listening on %s", addr)
		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			log.Fatal(err)
		}
	}()

	<-ctx.Done() // wait for Ctrl+C
	log.Println("shutting down")

	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := srv.Shutdown(shutdownCtx); err != nil {
		log.Println("shutdown:", err)
	}
}

curl -X POST localhost:8080/api/links -d '{"url":"https://go.dev/doc/"}'
curl -X POST localhost:8080/api/links -d '{"url":"https://pkg.go.dev/net/http"}'
curl -X POST localhost:8080/api/links -d '{"url":"go.dev"}'
curl -i localhost:8080/q0U
curl localhost:8080/api/links
curl localhost:8080/api/links/abc

// app_test.go
package main

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestEncodeBase62(t *testing.T) {
	tests := []struct {
		n    int
		want string
	}{
		{0, "0"}, {61, "Z"}, {62, "10"}, {100000, "q0U"},
	}
	for _, tc := range tests {
		if got := encodeBase62(tc.n); got != tc.want {
			t.Errorf("encodeBase62(%d) = %q, want %q", tc.n, got, tc.want)
		}
	}
}

func TestValidateURL(t *testing.T) {
	tests := []struct {
		in    string
		valid bool
	}{
		{"https://go.dev", true},
		{"http://example.in/menu?day=mon", true},
		{"go.dev", false},
		{"ftp://files.example.in", false},
		{"javascript:alert(1)", false},
		{"https://", false},
		{"", false},
	}
	for _, tc := range tests {
		err := validateURL(tc.in)
		if (err == nil) != tc.valid {
			t.Errorf("validateURL(%q) error = %v, want valid = %v", tc.in, err, tc.valid)
		}
	}
}

func TestCreateThenRedirect(t *testing.T) {
	app := &App{store: NewStore(), baseURL: "http://short.test"}
	h := app.routes()

	rec := httptest.NewRecorder()
	body := strings.NewReader(`{"url":"https://go.dev/doc/"}`)
	h.ServeHTTP(rec, httptest.NewRequest("POST", "/api/links", body))
	if rec.Code != http.StatusCreated {
		t.Fatalf("create: status %d, body %s", rec.Code, rec.Body)
	}
	var created linkJSON
	if err := json.NewDecoder(rec.Body).Decode(&created); err != nil {
		t.Fatal(err)
	}

	rec = httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest("GET", "/"+created.Code, nil))
	if rec.Code != http.StatusFound || rec.Header().Get("Location") != "https://go.dev/doc/" {
		t.Fatalf("redirect: status %d, location %q", rec.Code, rec.Header().Get("Location"))
	}

	if l, _ := app.store.Get(created.Code); l.Visits != 1 {
		t.Errorf("visits = %d, want 1", l.Visits)
	}
}

go test -race ./...