From 154c105345df6a474f5530cb21a16fb79f5998f6 Mon Sep 17 00:00:00 2001
From: Bob Lail <lail@squareup.com>
Date: Fri, 3 Jan 2025 13:19:47 -0800
Subject: [PATCH] feat: Parameterize the fallback handler to `httprule.Server`

Allow clients to define their own behavior in place of `http.NotFound` if none of the `httpRuleTemplates` were matched
---
 main.go                       | 3 ++-
 main_test.go                  | 2 +-
 serve/httprule/server.go      | 7 +++++--
 serve/httprule/server_test.go | 4 ++--
 4 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/main.go b/main.go
index b921c79..906d426 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"net/http"
 	"os"
 	"strings"
 
@@ -79,7 +80,7 @@ func (cs *cmdServe) Run(logLevel log.LogLevel) error {
 	}
 
 	if cs.HTTP {
-		h := httprule.NewServer(s.Files, s.UnknownHandler, logger, nil)
+		h := httprule.NewServer(s.Files, s.UnknownHandler, logger, nil, http.NotFoundHandler())
 		s.SetHTTPHandler(h)
 	}
 
diff --git a/main_test.go b/main_test.go
index 009c996..756f747 100644
--- a/main_test.go
+++ b/main_test.go
@@ -22,7 +22,7 @@ func TestHTTPRuleServer(t *testing.T) {
 	require.NoError(t, err)
 
 	ts := serve.NewUnstartedTestServer(serve.JsonnetEvaluator(), os.DirFS("serve/testdata/httpgreet"), opts...)
-	ts.SetHTTPHandler(httprule.NewServer(ts.Files, ts.UnknownHandler, log.DiscardLogger, nil))
+	ts.SetHTTPHandler(httprule.NewServer(ts.Files, ts.UnknownHandler, log.DiscardLogger, nil, http.NotFoundHandler()))
 	ts.Start()
 	defer ts.Stop()
 
diff --git a/serve/httprule/server.go b/serve/httprule/server.go
index 1a42840..39450d1 100644
--- a/serve/httprule/server.go
+++ b/serve/httprule/server.go
@@ -29,13 +29,15 @@ type Server struct {
 	httpMethods []*httpMethod
 	grpcHandler grpc.StreamHandler
 	log         log.Logger
+	next        http.Handler
 }
 
-func NewServer(files *registry.Files, handler grpc.StreamHandler, l log.Logger, httpRuleTemplates []*annotations.HttpRule) *Server {
+func NewServer(files *registry.Files, handler grpc.StreamHandler, l log.Logger, httpRuleTemplates []*annotations.HttpRule, next http.Handler) *Server {
 	return &Server{
 		httpMethods: loadHTTPRules(l, files, httpRuleTemplates),
 		grpcHandler: handler,
 		log:         l,
+		next:        next,
 	}
 }
 
@@ -46,7 +48,8 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 	}
-	http.NotFound(w, r)
+
+	s.next.ServeHTTP(w, r)
 }
 
 // Serve a google.api.http annotated method as HTTP
diff --git a/serve/httprule/server_test.go b/serve/httprule/server_test.go
index 3cda31b..d61e444 100644
--- a/serve/httprule/server_test.go
+++ b/serve/httprule/server_test.go
@@ -24,7 +24,7 @@ func TestHTTP(t *testing.T) {
 	ts := serve.NewTestServer(serve.JsonnetEvaluator(), os.DirFS("testdata/greet"), withLogger)
 	defer ts.Stop()
 
-	h := NewServer(ts.Files, ts.UnknownHandler, log.DiscardLogger, nil)
+	h := NewServer(ts.Files, ts.UnknownHandler, log.DiscardLogger, nil, http.NotFoundHandler())
 	ts.SetHTTPHandler(h)
 
 	body := `{"first_name": "Stranger"}`
@@ -102,7 +102,7 @@ func TestHTTPRuleInterpolation(t *testing.T) {
 		{Pattern: &annotations.HttpRule_Post{Post: "/post/{package}.{service}/{method}"}, Body: "*"},
 		{Pattern: &annotations.HttpRule_Get{Get: "/get/{method}"}},
 	}
-	h := NewServer(ts.Files, ts.UnknownHandler, logger, tmpl)
+	h := NewServer(ts.Files, ts.UnknownHandler, logger, tmpl, http.NotFoundHandler())
 	ts.SetHTTPHandler(h)
 
 	u := "http://" + ts.Addr() + "/get/SimpleHello"