Routing in Go Part 2: Upgrading the server so it behaves like a grown up

Part 2: Upgrading the server

Tests, auth middleware, a service layer, and request timeouts

In part 1 we built a small payment links API.

In part 2 we are going to do the things everyone promises they will do “after launch” and then never do until something catches fire.

We will add:

A service layer so handlers stop doing business decisions
Auth middleware with a simple API key
Request scoped timeouts so slow work does not hang forever
Proper handler tests with httptest

Same rule as before. No magic. Code that compiles.

Step 1: Restructure into a sane layout

Keep it small, but separate responsibilities.

Create this structure:

mkdir -p cmd/api internal/httpapi internal/links

You will end up with:

cmd/api/main.go bootstraps the server
internal/httpapi routing, middleware, helpers
internal/links store and service logic

Step 2: Move the domain and store into `internal/links`

Create internal/links/model.go

package links

import "time"

type PayLink struct {
	ID        int       `json:"id"`
	Amount    int       `json:"amount"`
	Currency  string    `json:"currency"`
	Note      string    `json:"note"`
	CreatedAt time.Time `json:"created_at"`
}

Create internal/links/store.go

package links

import (
	"sync"
	"time"
)

type Store interface {
	Create(amount int, currency, note string) PayLink
	List() []PayLink
	Get(id int) (PayLink, bool)
}

type MemoryStore struct {
	mu    sync.RWMutex
	next  int
	links map[int]PayLink
}

func NewMemoryStore() *MemoryStore {
	return &MemoryStore{
		next:  1,
		links: make(map[int]PayLink),
	}
}

func (s *MemoryStore) Create(amount int, currency, note string) PayLink {
	s.mu.Lock()
	defer s.mu.Unlock()

	l := PayLink{
		ID:        s.next,
		Amount:    amount,
		Currency:  currency,
		Note:      note,
		CreatedAt: time.Now().UTC(),
	}
	s.links[l.ID] = l
	s.next++
	return l
}

func (s *MemoryStore) List() []PayLink {
	s.mu.RLock()
	defer s.mu.RUnlock()

	out := make([]PayLink, 0, len(s.links))
	for _, l := range s.links {
		out = append(out, l)
	}
	return out
}

func (s *MemoryStore) Get(id int) (PayLink, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	l, ok := s.links[id]
	return l, ok
}

Step 3: Add a service layer that owns the rules

This is where you put “currency must be 3 letters” and “amount must be positive”.

Create internal/links/service.go

package links

import (
	"context"
	"errors"
	"strings"
	"time"
)

var (
	ErrInvalidAmount   = errors.New("amount must be greater than zero")
	ErrInvalidCurrency = errors.New("currency must be a 3 letter code like USD")
	ErrNotFound        = errors.New("link not found")
)

type Service struct {
	store Store
}

func NewService(store Store) *Service {
	return &Service{store: store}
}

func (s *Service) Create(ctx context.Context, amount int, currency, note string) (PayLink, error) {
	select {
	case <-ctx.Done():
		return PayLink{}, ctx.Err()
	default:
	}

	currency = strings.ToUpper(strings.TrimSpace(currency))
	note = strings.TrimSpace(note)

	if amount <= 0 {
		return PayLink{}, ErrInvalidAmount
	}
	if len(currency) != 3 {
		return PayLink{}, ErrInvalidCurrency
	}

	// Pretend we might do more work here later.
	time.Sleep(5 * time.Millisecond)

	return s.store.Create(amount, currency, note), nil
}

func (s *Service) List(ctx context.Context) ([]PayLink, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
	}
	return s.store.List(), nil
}

func (s *Service) Get(ctx context.Context, id int) (PayLink, error) {
	select {
	case <-ctx.Done():
		return PayLink{}, ctx.Err()
	default:
	}
	l, ok := s.store.Get(id)
	if !ok {
		return PayLink{}, ErrNotFound
	}
	return l, nil
}

Notice what happened.

Handlers no longer need to know the rules. They only translate HTTP into method calls.

That is how you stop your routing from turning into a soap opera.

Step 4: HTTP helpers and middleware

Create internal/httpapi/helpers.go

package httpapi

import (
	"encoding/json"
	"net/http"
)

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

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

Create internal/httpapi/middleware.go

package httpapi

import (
	"log"
	"math/rand"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
)

type Middleware func(http.Handler) http.Handler

func Chain(h http.Handler, m ...Middleware) http.Handler {
	for i := len(m) - 1; i >= 0; i-- {
		h = m[i](h)
	}
	return h
}

func RequestID(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		id := r.Header.Get("X-Request-Id")
		if id == "" {
			id = newRequestID()
		}
		w.Header().Set("X-Request-Id", id)
		next.ServeHTTP(w, r)
	})
}

func Logging(logger *log.Logger) Middleware {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()
			next.ServeHTTP(w, r)
			logger.Printf("%s %s from=%s dur=%s",
				r.Method,
				r.URL.Path,
				clientIP(r),
				time.Since(start).Round(time.Millisecond),
			)
		})
	}
}

func APIKeyAuth(expected string) Middleware {
	expected = strings.TrimSpace(expected)
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Let health stay public. Your uptime monitor is not your enemy.
			if r.URL.Path == "/health" {
				next.ServeHTTP(w, r)
				return
			}

			if expected == "" {
				// If no key is configured, do not pretend you are secure.
				WriteError(w, http.StatusInternalServerError, "server missing API key configuration")
				return
			}

			got := strings.TrimSpace(r.Header.Get("X-API-Key"))
			if got != expected {
				WriteError(w, http.StatusUnauthorized, "unauthorized")
				return
			}

			next.ServeHTTP(w, r)
		})
	}
}

func clientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

func newRequestID() string {
	return strconv.FormatInt(time.Now().UnixNano(), 36) + "-" + strconv.Itoa(rand.Intn(1000000))
}

Step 5: Handlers that use the service layer

Create internal/httpapi/handlers.go

package httpapi

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"strconv"
	"time"

	"example.com/paylinks/internal/links"
)

type API struct {
	Links *links.Service
}

func (a *API) Routes() http.Handler {
	mux := http.NewServeMux()

	mux.HandleFunc("GET /health", func(w http.ResponseWriter, r *http.Request) {
		WriteJSON(w, http.StatusOK, map[string]string{"status": "ok"})
	})

	mux.HandleFunc("/links", a.linksCollection)
	mux.HandleFunc("GET /links/{id}", a.linksGet)

	return mux
}

func (a *API) linksCollection(w http.ResponseWriter, r *http.Request) {
	ctx, cancel := context.WithTimeout(r.Context(), 2*time.Second)
	defer cancel()

	switch r.Method {
	case http.MethodGet:
		items, err := a.Links.List(ctx)
		if err != nil {
			WriteError(w, http.StatusRequestTimeout, "request timed out")
			return
		}
		WriteJSON(w, http.StatusOK, items)
		return

	case http.MethodPost:
		var body struct {
			Amount   int    `json:"amount"`
			Currency string `json:"currency"`
			Note     string `json:"note"`
		}
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			WriteError(w, http.StatusBadRequest, "invalid JSON body")
			return
		}

		l, err := a.Links.Create(ctx, body.Amount, body.Currency, body.Note)
		if err != nil {
			switch {
			case errors.Is(err, links.ErrInvalidAmount), errors.Is(err, links.ErrInvalidCurrency):
				WriteError(w, http.StatusBadRequest, err.Error())
			case errors.Is(err, context.DeadlineExceeded):
				WriteError(w, http.StatusRequestTimeout, "request timed out")
			default:
				WriteError(w, http.StatusInternalServerError, "server error")
			}
			return
		}

		WriteJSON(w, http.StatusCreated, l)
		return

	default:
		WriteError(w, http.StatusMethodNotAllowed, "method not allowed")
		return
	}
}

func (a *API) linksGet(w http.ResponseWriter, r *http.Request) {
	ctx, cancel := context.WithTimeout(r.Context(), 2*time.Second)
	defer cancel()

	idStr := r.PathValue("id")
	id, err := strconv.Atoi(idStr)
	if err != nil || id <= 0 {
		WriteError(w, http.StatusBadRequest, "invalid link id")
		return
	}

	l, err := a.Links.Get(ctx, id)
	if err != nil {
		switch {
		case errors.Is(err, links.ErrNotFound):
			WriteError(w, http.StatusNotFound, "link not found")
		case errors.Is(err, context.DeadlineExceeded):
			WriteError(w, http.StatusRequestTimeout, "request timed out")
		default:
			WriteError(w, http.StatusInternalServerError, "server error")
		}
		return
	}

	WriteJSON(w, http.StatusOK, l)
}

Step 6: Wire it all together in `cmd/api/main.go`

Create cmd/api/main.go

package main

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

	"example.com/paylinks/internal/httpapi"
	"example.com/paylinks/internal/links"
)

func main() {
	logger := log.New(os.Stdout, "", log.LstdFlags)

	apiKey := os.Getenv("API_KEY")
	store := links.NewMemoryStore()
	service := links.NewService(store)

	api := &httpapi.API{Links: service}
	routes := api.Routes()

	handler := httpapi.Chain(
		routes,
		httpapi.RequestID,
		httpapi.Logging(logger),
		httpapi.APIKeyAuth(apiKey),
	)

	srv := &http.Server{
		Addr:              ":8080",
		Handler:           handler,
		ReadHeaderTimeout: 5 * time.Second,
		ReadTimeout:       10 * time.Second,
		WriteTimeout:      10 * time.Second,
		IdleTimeout:       60 * time.Second,
	}

	go func() {
		logger.Println("listening on http://localhost:8080")
		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			logger.Fatalf("server error: %v", err)
		}
	}()

	waitForShutdown(logger, srv)
}

func waitForShutdown(logger *log.Logger, srv *http.Server) {
	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)

	<-stop
	logger.Println("shutting down...")

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if err := srv.Shutdown(ctx); err != nil {
		logger.Printf("shutdown error: %v", err)
	}
	logger.Println("bye")
}

Run it:

API_KEY=supersecret go run ./cmd/api

Try without auth:

curl -s http://localhost:8080/links

You should get unauthorized.

Try with auth:

curl -s http://localhost:8080/links -H "X-API-Key: supersecret"

Now it works.

This is the exact point where a lot of people tell themselves they have “security”. You do not. You have a locked door. It is still a wooden door. But it is better than a sign that says “please do not hack”.

Step 7: Write tests that prove your routes behave

This is where confidence comes from. Not vibes.

Create internal/httpapi/handlers_test.go

package httpapi

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"testing"

	"example.com/paylinks/internal/links"
)

func newTestServer() http.Handler {
	store := links.NewMemoryStore()
	service := links.NewService(store)
	api := &API{Links: service}
	return api.Routes()
}

func TestHealth(t *testing.T) {
	h := newTestServer()

	req := httptest.NewRequest(http.MethodGet, "/health", nil)
	rec := httptest.NewRecorder()

	h.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", rec.Code)
	}
}

func TestCreateAndGetLink(t *testing.T) {
	h := newTestServer()

	// Create
	body := []byte(`{"amount":5000,"currency":"kes","note":"invoice 183"}`)
	req := httptest.NewRequest(http.MethodPost, "/links", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()

	h.ServeHTTP(rec, req)

	if rec.Code != http.StatusCreated {
		t.Fatalf("expected 201, got %d body=%s", rec.Code, rec.Body.String())
	}

	// Get
	req2 := httptest.NewRequest(http.MethodGet, "/links/1", nil)
	rec2 := httptest.NewRecorder()

	h.ServeHTTP(rec2, req2)

	if rec2.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d body=%s", rec2.Code, rec2.Body.String())
	}
}

func TestCreateValidation(t *testing.T) {
	h := newTestServer()

	body := []byte(`{"amount":0,"currency":"usd"}`)
	req := httptest.NewRequest(http.MethodPost, "/links", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()

	h.ServeHTTP(rec, req)

	if rec.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d body=%s", rec.Code, rec.Body.String())
	}
}

Run tests:

go test ./...

If it passes, congrats. You now have a codebase that can be changed without praying.

Step 8: Test auth middleware without hating your life

Right now tests call api.Routes() directly, so they bypass middleware.

That is good. You should test handlers and middleware separately.

Create internal/httpapi/middleware_test.go

package httpapi

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestAPIKeyAuth(t *testing.T) {
	final := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	h := APIKeyAuth("secret")(final)

	// Missing key
	req := httptest.NewRequest(http.MethodGet, "/links", nil)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)

	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("expected 401, got %d", rec.Code)
	}

	// Correct key
	req2 := httptest.NewRequest(http.MethodGet, "/links", nil)
	req2.Header.Set("X-API-Key", "secret")
	rec2 := httptest.NewRecorder()
	h.ServeHTTP(rec2, req2)

	if rec2.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", rec2.Code)
	}

	// Health is public
	req3 := httptest.NewRequest(http.MethodGet, "/health", nil)
	rec3 := httptest.NewRecorder()
	h.ServeHTTP(rec3, req3)

	if rec3.Code != http.StatusOK {
		t.Fatalf("expected 200 for health, got %d", rec3.Code)
	}
}

Run:

go test ./...

What we fixed, in human terms

You now have:

Routing that stays clean because handlers are thin
Business rules that live in a service, not scattered around
Middleware that centralizes behavior like auth
Timeouts so requests do not hang forever
Tests that turn fear into evidence

Also, you avoided the classic trap where the router becomes your entire application.

Your router is now just traffic control. As it should be.

Share this writing

Twitter LinkedIn

More notes

Routing in Go Part 1: A small server you can grow without rewriting at midnight

Routing in Go Part 3: The last upgrade