main.go

// Copyright 2024 The Tessera authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// aws is a simple personality allowing to run conformance/compliance/performance tests and showing how to use the Tessera AWS storage implementation.
package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	"log/slog"

	aaws "github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/go-sql-driver/mysql"
	fnote "github.com/transparency-dev/formats/note"
	"github.com/transparency-dev/tessera"
	"github.com/transparency-dev/tessera/storage/aws"
	aws_as "github.com/transparency-dev/tessera/storage/aws/antispam"
	"golang.org/x/mod/sumdb/note"
)

var (
	bucket            = flag.String("bucket", "", "Bucket to use for storing log")
	dbName            = flag.String("db_name", "", "AuroraDB name for the log DB")
	dbHost            = flag.String("db_host", "", "AuroraDB host")
	dbPort            = flag.Int("db_port", 3306, "AuroraDB port")
	dbUser            = flag.String("db_user", "", "AuroraDB user")
	dbPassword        = flag.String("db_password", "", "AuroraDB user")
	dbMaxConns        = flag.Int("db_max_conns", 0, "Maximum connections to the database, defaults to 0, i.e unlimited")
	dbMaxIdle         = flag.Int("db_max_idle_conns", 2, "Maximum idle database connections in the connection pool, defaults to 2")
	s3Endpoint        = flag.String("s3_endpoint", "", "Endpoint for custom non-AWS S3 service")
	s3AccessKeyID     = flag.String("s3_access_key", "", "Access key ID for custom non-AWS S3 service")
	s3SecretAccessKey = flag.String("s3_secret", "", "Secret access key for custom non-AWS S3 service")

	listen            = flag.String("listen", ":2024", "Address:port to listen on")
	signer            = flag.String("signer", "", "Note signer to use to sign checkpoints")
	publishInterval   = flag.Duration("publish_interval", 3*time.Second, "How frequently to publish updated checkpoints")
	traceFraction     = flag.Float64("trace_fraction", 0, "Fraction of open-telemetry span traces to sample")
	slogLevel         = flag.Int("slog_level", 0, "The cut-off threshold for structured logging. Default is 0 (INFO). See https://pkg.go.dev/log/slog#Level for other levels.")
	logFormat         = flag.String("log_format", "text", "The format of the logs: text or json.")
	additionalSigners = []string{}

	antispamEnable = flag.Bool("antispam", false, "EXPERIMENTAL: Set to true to enable persistent antispam storage")
	antispamDb     = flag.String("antispam_db_name", "", "AuroraDB name for the antispam DB")
)

func init() {
	flag.Func("additional_signer", "Additional note signer for checkpoints, may be specified multiple times", func(s string) error {
		additionalSigners = append(additionalSigners, s)
		return nil
	})
}

func main() {
	flag.Parse()
	ctx := context.Background()
	var handler slog.Handler
	switch *logFormat {
	case "json":
		handler = slog.NewJSONHandler(os.Stderr, &slog.HandlerOptions{Level: slog.Level(*slogLevel)})
	default:
		handler = slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.Level(*slogLevel)})
	}
	slog.SetDefault(slog.New(handler))

	shutdownOTel := initOTel(ctx, *traceFraction)
	defer shutdownOTel(ctx)
	s, a := signerFromFlags()

	// Create our Tessera storage backend:
	awsCfg := storageConfigFromFlags()
	driver, err := aws.New(ctx, awsCfg)
	if err != nil {
		slog.ErrorContext(ctx, "Failed to create new AWS storage", slog.Any("error", err))
		os.Exit(1)
	}
	var antispam tessera.Antispam
	// Persistent antispam is currently experimental, so there's no documentation yet!
	if *antispamEnable {
		asOpts := aws_as.AntispamOpts{} // Use defaults
		antispam, err = aws_as.NewAntispam(ctx, antispamMysqlConfig().FormatDSN(), asOpts)
		if err != nil {
			slog.ErrorContext(ctx, "Failed to create new AWS antispam storage", slog.Any("error", err))
			os.Exit(1)
		}
	}
	appender, shutdown, _, err := tessera.NewAppender(ctx, driver, tessera.NewAppendOptions().
		WithCheckpointSigner(s, a...).
		WithCheckpointInterval(*publishInterval).
		WithBatching(512, 300*time.Millisecond).
		WithPushback(10*4096).
		WithAntispam(tessera.DefaultAntispamInMemorySize, antispam))
	if err != nil {
		slog.ErrorContext(ctx, "Failed to create new appender", slog.Any("error", err))
		os.Exit(1)
	}

	// Expose a HTTP handler for the conformance test writes.
	// This should accept arbitrary bytes POSTed to /add, and return an ascii
	// decimal representation of the index assigned to the entry.
	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
		}

		idx, err := appender.Add(r.Context(), tessera.NewEntry(b))()
		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
		}
		// Write out the assigned index
		_, _ = fmt.Fprintf(w, "%d", idx.Index)
	})

	var protocols http.Protocols
	protocols.SetHTTP1(true)
	protocols.SetUnencryptedHTTP2(true)

	server := &http.Server{
		Addr:              *listen,
		Handler:           http.DefaultServeMux,
		Protocols:         &protocols,
		ReadHeaderTimeout: 5 * time.Second,
	}

	if err := server.ListenAndServe(); err != nil {
		if err := shutdown(ctx); err != nil {
			slog.ErrorContext(ctx, "Failed to cleanly shutdown after ListenAndServe", slog.Any("error", err))
			os.Exit(1)
		}
		slog.ErrorContext(ctx, "ListenAndServe", slog.Any("error", err))
		os.Exit(1)
	}
}

// storageConfigFromFlags returns an aws.Config struct populated with values
// provided via flags.
func storageConfigFromFlags() aws.Config {
	ctx := context.Background()
	if *bucket == "" {
		slog.ErrorContext(ctx, "--bucket must be set")
		os.Exit(1)
	}
	if *dbName == "" {
		slog.ErrorContext(ctx, "--db_name must be set")
		os.Exit(1)
	}
	if *dbHost == "" {
		slog.ErrorContext(ctx, "--db_host must be set")
		os.Exit(1)
	}
	if *dbPort == 0 {
		slog.ErrorContext(ctx, "--db_port must be set")
		os.Exit(1)
	}
	if *dbUser == "" {
		slog.ErrorContext(ctx, "--db_user must be set")
		os.Exit(1)
	}
	// Empty password isn't an option with AuroraDB MySQL.
	if *dbPassword == "" {
		slog.ErrorContext(ctx, "--db_password must be set")
		os.Exit(1)
	}

	c := mysql.Config{
		User:                    *dbUser,
		Passwd:                  *dbPassword,
		Net:                     "tcp",
		Addr:                    fmt.Sprintf("%s:%d", *dbHost, *dbPort),
		DBName:                  *dbName,
		AllowCleartextPasswords: true,
		AllowNativePasswords:    true,
	}

	// Configure to use MinIO Server
	var awsConfig *aaws.Config
	var s3Opts func(o *s3.Options)
	if *s3Endpoint != "" {
		const defaultRegion = "us-east-1"
		s3Opts = func(o *s3.Options) {
			o.BaseEndpoint = aaws.String(*s3Endpoint)
			o.Credentials = credentials.NewStaticCredentialsProvider(*s3AccessKeyID, *s3SecretAccessKey, "")
			o.Region = defaultRegion
			o.UsePathStyle = true
		}

		awsConfig = &aaws.Config{
			Region: defaultRegion,
		}
	}

	return aws.Config{
		Bucket:       *bucket,
		SDKConfig:    awsConfig,
		S3Options:    s3Opts,
		DSN:          c.FormatDSN(),
		MaxOpenConns: *dbMaxConns,
		MaxIdleConns: *dbMaxIdle,
	}
}

func antispamMysqlConfig() *mysql.Config {
	ctx := context.Background()
	if *antispamDb == "" {
		slog.ErrorContext(ctx, "--antispam_db_name must be set")
		os.Exit(1)
	}
	if *dbHost == "" {
		slog.ErrorContext(ctx, "--db_host must be set")
		os.Exit(1)
	}
	if *dbPort == 0 {
		slog.ErrorContext(ctx, "--db_port must be set")
		os.Exit(1)
	}
	if *dbUser == "" {
		slog.ErrorContext(ctx, "--db_user must be set")
		os.Exit(1)
	}
	// Empty password isn't an option with AuroraDB MySQL.
	if *dbPassword == "" {
		slog.ErrorContext(ctx, "--db_password must be set")
		os.Exit(1)
	}

	return &mysql.Config{
		User:                    *dbUser,
		Passwd:                  *dbPassword,
		Net:                     "tcp",
		Addr:                    fmt.Sprintf("%s:%d", *dbHost, *dbPort),
		DBName:                  *antispamDb,
		AllowCleartextPasswords: true,
		AllowNativePasswords:    true,
	}
}

func signerFromFlags() (note.Signer, []note.Signer) {
	s, err := fnote.NewSigner(*signer)
	if err != nil {
		slog.ErrorContext(context.Background(), "Failed to create new signer", slog.Any("error", err))
		os.Exit(1)
	}

	var a []note.Signer
	for _, as := range additionalSigners {
		s, err := fnote.NewSigner(as)
		if err != nil {
			slog.ErrorContext(context.Background(), "Failed to create additional signer", slog.Any("error", err))
			os.Exit(1)
		}
		a = append(a, s)
	}

	return s, a
}