Sample personality library proof-of-concept

transparency-dev · Aug 20, 2024 · a5d575c · a5d575c
1 parent 38226b6
commit a5d575c
Show file tree

Hide file tree

Showing 3 changed files with 122 additions and 65 deletions.
diff --git a/cmd/example-gcp/main.go b/cmd/example-gcp/main.go
@@ -30,6 +30,7 @@ import (
 
 	kms "cloud.google.com/go/kms/apiv1"
 	tessera "github.com/transparency-dev/trillian-tessera"
+	samplepersonality "github.com/transparency-dev/trillian-tessera/personalities/sample"
 	"github.com/transparency-dev/trillian-tessera/storage/gcp"
 	"golang.org/x/mod/sumdb/note"
 	"k8s.io/klog/v2"
@@ -74,29 +75,6 @@ func main() {
 		klog.Exitf("Failed to create new GCP storage: %v", err)
 	}
 
-	http.HandleFunc("POST /add", func(w http.ResponseWriter, r *http.Request) {
-		b, err := io.ReadAll(r.Body)
-		if err != nil {
-			w.WriteHeader(http.StatusInternalServerError)
-			return
-		}
-		defer r.Body.Close()
-
-		id := sha256.Sum256(b)
-		idx, err := storage.Add(r.Context(), tessera.NewEntry(b, tessera.WithIdentity(id[:])))
-		if err != nil {
-			if errors.Is(err, tessera.ErrPushback) {
-				w.Header().Add("Retry-After", "1")
-				w.WriteHeader(http.StatusServiceUnavailable)
-				return
-			}
-			w.WriteHeader(http.StatusInternalServerError)
-			_, _ = w.Write([]byte(err.Error()))
-			return
-		}
-		_, _ = w.Write([]byte(fmt.Sprintf("%d", idx)))
-	})
-
 	// TODO: remove this proxy
 	serveGCS := func(w http.ResponseWriter, r *http.Request) {
 		resource := strings.TrimLeft(r.URL.Path, "/")
@@ -109,12 +87,14 @@ func main() {
 		}
 		_, _ = w.Write(b)
 	}
-	http.HandleFunc("GET /checkpoint", serveGCS)
-	http.HandleFunc("GET /tile/", serveGCS)
 
-	if err := http.ListenAndServe(*listen, http.DefaultServeMux); err != nil {
-		klog.Exitf("ListenAndServe: %v", err)
-	}
+	personality := samplepersonality.New(ctx, samplepersonality.Handlers{
+		Checkpoint:  serveGCS,
+		Tile:        serveGCS,
+		EntryBundle: serveGCS,
+		Add:         addHandler(storage),
+	})
+	personality.Run(*listen)
 }
 
 // signerFromFlags creates and returns a new KMSSigner from the flags, along with a close func.
@@ -138,3 +118,28 @@ func signerFromFlags(ctx context.Context) (note.Signer, note.Verifier, func() er
 
 	return signer, verifier, kmClient.Close
 }
+
+func addHandler(storage *gcp.Storage) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		b, err := io.ReadAll(r.Body)
+		if err != nil {
+			w.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+		defer r.Body.Close()
+
+		id := sha256.Sum256(b)
+		idx, err := storage.Add(r.Context(), tessera.NewEntry(b, tessera.WithIdentity(id[:])))
+		if err != nil {
+			if errors.Is(err, tessera.ErrPushback) {
+				w.Header().Add("Retry-After", "1")
+				w.WriteHeader(http.StatusServiceUnavailable)
+				return
+			}
+			w.WriteHeader(http.StatusInternalServerError)
+			_, _ = w.Write([]byte(err.Error()))
+			return
+		}
+		_, _ = w.Write([]byte(fmt.Sprintf("%d", idx)))
+	}
+}
diff --git a/cmd/example-mysql/main.go b/cmd/example-mysql/main.go
@@ -27,6 +27,7 @@ import (
 
 	tessera "github.com/transparency-dev/trillian-tessera"
 	"github.com/transparency-dev/trillian-tessera/api/layout"
+	samplepersonality "github.com/transparency-dev/trillian-tessera/personalities/sample"
 	"github.com/transparency-dev/trillian-tessera/storage/mysql"
 	"golang.org/x/mod/sumdb/note"
 	"k8s.io/klog/v2"
@@ -80,7 +81,43 @@ func main() {
 		klog.Exitf("Failed to create new MySQL storage: %v", err)
 	}
 
-	http.HandleFunc("GET /checkpoint", func(w http.ResponseWriter, r *http.Request) {
+	personality := samplepersonality.New(ctx, samplepersonality.Handlers{
+		Checkpoint:  checkpointHandler(storage),
+		Tile:        tileHandler(storage),
+		EntryBundle: entryBundleHandler(storage),
+		Add:         addHandler(storage),
+	})
+	personality.Run(*listen)
+}
+
+func initDatabaseSchema(ctx context.Context) {
+	if *initSchemaPath != "" {
+		klog.Infof("Initializing database schema")
+
+		db, err := sql.Open("mysql", *mysqlURI+"?multiStatements=true")
+		if err != nil {
+			klog.Exitf("Failed to connect to DB: %v", err)
+		}
+		defer func() {
+			if err := db.Close(); err != nil {
+				klog.Warningf("Failed to close db: %v", err)
+			}
+		}()
+
+		rawSchema, err := os.ReadFile(*initSchemaPath)
+		if err != nil {
+			klog.Exitf("Failed to read init schema file %q: %v", *initSchemaPath, err)
+		}
+		if _, err := db.ExecContext(ctx, string(rawSchema)); err != nil {
+			klog.Exitf("Failed to execute init database schema: %v", err)
+		}
+
+		klog.Infof("Database schema initialized")
+	}
+}
+
+func checkpointHandler(storage *mysql.Storage) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
 		checkpoint, err := storage.ReadCheckpoint(r.Context())
 		if err != nil {
 			klog.Errorf("/checkpoint: %v", err)
@@ -96,9 +133,11 @@ func main() {
 			klog.Errorf("/checkpoint: %v", err)
 			return
 		}
-	})
+	}
+}
 
-	http.HandleFunc("GET /tile/{level}/{index...}", func(w http.ResponseWriter, r *http.Request) {
+func tileHandler(storage *mysql.Storage) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
 		level, index, width, err := layout.ParseTileLevelIndexWidth(r.PathValue("level"), r.PathValue("index"))
 		if err != nil {
 			w.WriteHeader(http.StatusBadRequest)
@@ -125,9 +164,11 @@ func main() {
 			klog.Errorf("/tile/{level}/{index...}: %v", err)
 			return
 		}
-	})
+	}
+}
 
-	http.HandleFunc("GET /tile/entries/{index...}", func(w http.ResponseWriter, r *http.Request) {
+func entryBundleHandler(storage *mysql.Storage) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
 		index, _, err := layout.ParseTileIndexWidth(r.PathValue("index"))
 		if err != nil {
 			w.WriteHeader(http.StatusBadRequest)
@@ -154,9 +195,11 @@ func main() {
 			klog.Errorf("/tile/entries/{index...}: %v", err)
 			return
 		}
-	})
+	}
+}
 
-	http.HandleFunc("POST /add", func(w http.ResponseWriter, r *http.Request) {
+func addHandler(storage *mysql.Storage) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
 		b, err := io.ReadAll(r.Body)
 		if err != nil {
 			w.WriteHeader(http.StatusInternalServerError)
@@ -177,35 +220,5 @@ func main() {
 			klog.Errorf("/add: %v", err)
 			return
 		}
-	})
-
-	if err := http.ListenAndServe(*listen, http.DefaultServeMux); err != nil {
-		klog.Exitf("ListenAndServe: %v", err)
-	}
-}
-
-func initDatabaseSchema(ctx context.Context) {
-	if *initSchemaPath != "" {
-		klog.Infof("Initializing database schema")
-
-		db, err := sql.Open("mysql", *mysqlURI+"?multiStatements=true")
-		if err != nil {
-			klog.Exitf("Failed to connect to DB: %v", err)
-		}
-		defer func() {
-			if err := db.Close(); err != nil {
-				klog.Warningf("Failed to close db: %v", err)
-			}
-		}()
-
-		rawSchema, err := os.ReadFile(*initSchemaPath)
-		if err != nil {
-			klog.Exitf("Failed to read init schema file %q: %v", *initSchemaPath, err)
-		}
-		if _, err := db.ExecContext(ctx, string(rawSchema)); err != nil {
-			klog.Exitf("Failed to execute init database schema: %v", err)
-		}
-
-		klog.Infof("Database schema initialized")
 	}
 }
diff --git a/personalities/sample/sample.go b/personalities/sample/sample.go
@@ -0,0 +1,39 @@
+package samplepersonality
+
+import (
+	"context"
+	"net/http"
+
+	"k8s.io/klog/v2"
+)
+
+type SamplePersonality struct {
+	handlers Handlers
+}
+
+type Handlers struct {
+	// Read
+	Checkpoint,
+	Tile,
+	EntryBundle,
+
+	// Write
+	Add func(http.ResponseWriter, *http.Request)
+}
+
+func New(ctx context.Context, handlers Handlers) SamplePersonality {
+	return SamplePersonality{
+		handlers: handlers,
+	}
+}
+
+func (p *SamplePersonality) Run(addr string) {
+	http.HandleFunc("GET /checkpoint", p.handlers.Checkpoint)
+	http.HandleFunc("GET /tile/{level}/{index...}", p.handlers.Tile)
+	http.HandleFunc("GET /tile/entries/{index...}", p.handlers.EntryBundle)
+	http.HandleFunc("POST /add", p.handlers.Add)
+
+	if err := http.ListenAndServe(addr, http.DefaultServeMux); err != nil {
+		klog.Exitf("ListenAndServe: %v", err)
+	}
+}