proxy TTL

Author: walterwanderley

.gitignore

@@ -1,2 +1,5 @@
 *.so
 *.db
+*.db-shm
+*.db-wal
+

cmd/sqlite-http-proxy/main.go

@@ -4,39 +4,45 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"database/sql"
-	"errors"
 	"flag"
 	"fmt"
 	"log"
 	"net"
 	"net/http"
 	"os"
 	"strings"
-	"time"
 
 	"github.com/elazarl/goproxy"
 	_ "github.com/mattn/go-sqlite3"
 
 	"github.com/walterwanderley/sqlite-http-cache/db"
 )
 
-func main() {
-	var (
-		port       uint
-		allowHTTP2 bool
-		verbose    bool
+var (
+	port       uint
+	allowHTTP2 bool
+	verbose    bool
+
+	ttl uint
 
-		caCert    string
-		caCertKey string
+	caCert    string
+	caCertKey string
+
+	responseTables    string
+	forceCreateTables bool
+	readOnly          bool
+)
 
-		responseTables string
-	)
+func main() {
 	flag.UintVar(&port, "p", 8080, "Server port")
 	flag.BoolVar(&verbose, "v", false, "Enable verbose mode")
 	flag.BoolVar(&allowHTTP2, "h2", false, "Allow HTTP2")
+	flag.UintVar(&ttl, "ttl", 0, "Time to Live in seconds (0 is infinite time)")
 	flag.StringVar(&responseTables, "response-tables", "", "Comma separated list of database tables used to store response data")
+	flag.BoolVar(&forceCreateTables, "force-create-tables", false, "Force create response tables if not exists")
 	flag.StringVar(&caCert, "ca-cert", "", "Path to CA Certificate file (required to HTTPS proxy)")
 	flag.StringVar(&caCertKey, "ca-cert-key", "", "Path to CA Certificate Key file (required to HTTPS proxy)")
+	flag.BoolVar(&readOnly, "ro", false, "Read Only mode. Do not store new HTTP responses")
 	flag.Parse()
 
 	if len(flag.Args()) != 1 {
@@ -58,6 +64,12 @@ func main() {
 		}
 	} else {
 		tableList = strings.Split(responseTables, ",")
+		if forceCreateTables {
+			err := db.CreateResponseTables(sqlDB, tableList...)
+			if err != nil {
+				log.Fatalf("force create tables: %v", err)
+			}
+		}
 	}
 
 	repository, err := db.NewRepository(sqlDB, tableList...)
@@ -86,30 +98,14 @@ func main() {
 		proxy.Logger.Printf("INFO: Starting HTTP Proxy...")
 	}
 
-	proxy.OnRequest().DoFunc(
-		func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
-			if r.Method != http.MethodGet || ctx.Req == nil || ctx.Req.URL == nil {
-				return r, nil
-			}
-
-			url := ctx.Req.URL.String()
-			resp, err := repository.FindByURL(r.Context(), url)
-			if err != nil {
-				if !errors.Is(err, sql.ErrNoRows) {
-					proxy.Logger.Printf("ERROR: query error: %s", err.Error())
-				}
-				return r, nil
-			}
-			if verbose {
-				proxy.Logger.Printf("INFO: serving from database url=%s status=%d timestamp=%s", url, resp.Status, resp.Timestamp.Format(time.RFC3339))
-			}
-
-			return r, &http.Response{
-				StatusCode: resp.Status,
-				Body:       resp.Body,
-				Header:     http.Header(resp.Headers),
-			}
+	proxy.OnRequest().Do(&requestHandler{
+		querier: repository,
+	})
+	if !readOnly {
+		proxy.OnResponse().Do(&responseHandler{
+			writer: repository,
 		})
+	}
 
 	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
 	if err != nil {

cmd/sqlite-http-proxy/request.go

@@ -0,0 +1,54 @@
+package main
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"log/slog"
+	"net/http"
+	"time"
+
+	"github.com/elazarl/goproxy"
+
+	"github.com/walterwanderley/sqlite-http-cache/db"
+)
+
+type requestQuerier interface {
+	FindByURL(ctx context.Context, url string) (*db.Response, error)
+}
+
+type requestHandler struct {
+	querier requestQuerier
+}
+
+func (h *requestHandler) Handle(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
+	if r.Method != http.MethodGet {
+		return r, nil
+	}
+
+	url := ctx.Req.URL.String()
+	resp, err := h.querier.FindByURL(r.Context(), url)
+	if err != nil {
+		if !errors.Is(err, sql.ErrNoRows) {
+			slog.Error("database query", "error", err.Error())
+		}
+		// tell the responseHandler to save the new response data
+		ctx.UserData = ""
+		return r, nil
+	}
+
+	if !readOnly && uint(time.Since(resp.Timestamp).Seconds()) > ttl {
+		// data is too old, tell the responseHandler to save the new data
+		ctx.UserData = resp.TableName
+		return r, nil
+	}
+	if verbose {
+		slog.Info("serving from database", "url", url, "status", resp.Status, "timestamp", resp.Timestamp.Format(time.RFC3339))
+	}
+
+	return r, &http.Response{
+		StatusCode: resp.Status,
+		Body:       resp.Body,
+		Header:     http.Header(resp.Headers),
+	}
+}

cmd/sqlite-http-proxy/response.go

@@ -0,0 +1,37 @@
+package main
+
+import (
+	"context"
+	"log/slog"
+	"net/http"
+
+	"github.com/elazarl/goproxy"
+	"github.com/walterwanderley/sqlite-http-cache/db"
+)
+
+type responseWriter interface {
+	Write(ctx context.Context, url string, resp *db.Response) error
+}
+
+type responseHandler struct {
+	writer responseWriter
+}
+
+func (h *responseHandler) Handle(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
+	if ctx.UserData != nil {
+		if verbose {
+			slog.Info("recording response", "url", ctx.Req.URL.String(), "status", resp.StatusCode)
+		}
+		responseDB, err := db.HttpToResponse(resp)
+		if err != nil {
+			slog.Error("adapter response body", "error", err)
+		} else {
+			responseDB.TableName = ctx.UserData.(string)
+			err := h.writer.Write(context.Background(), ctx.Req.URL.String(), responseDB)
+			if err != nil {
+				slog.Error("recording response", "error", err, "url", ctx.Req.URL.String(), "status", resp.StatusCode)
+			}
+		}
+	}
+	return resp
+}

cmd/sqlite-http-refresh/main.go

@@ -15,6 +15,7 @@ import (
 	_ "github.com/mattn/go-sqlite3"
 
 	"github.com/walterwanderley/sqlite-http-cache/config"
+	"github.com/walterwanderley/sqlite-http-cache/db"
 	_ "github.com/walterwanderley/sqlite-http-cache/extension"
 )
 
@@ -26,7 +27,7 @@ var (
 	timeout           uint
 	insecure          bool
 	ignoreStatusError bool
-	responseTableName string
+	responseTables    string
 
 	oauth2ClientID     string
 	oauth2ClientSecret string
@@ -41,12 +42,12 @@ func main() {
 	// Scheduler strategy
 	flag.DurationVar(&interval, "check-interval", 30*time.Second, "Interval to wait for check expired data")
 	flag.UintVar(&ttl, "ttl", 30*60, "Time to Live in seconds")
-	flag.StringVar(&matchURL, "match-url", "", "Filter URLs")
+	flag.StringVar(&matchURL, "match-url", "%", "Filter URLs (SQL syntax)")
 	// Request/Store config
 	flag.UintVar(&timeout, "timeout", 30*1000, "Timeout in milliseconds")
 	flag.BoolVar(&insecure, "insecure", false, "Disable TLS verification")
 	flag.BoolVar(&ignoreStatusError, "ignore-status-error", false, "ignore responses with status code != 2xx")
-	flag.StringVar(&responseTableName, "response-table", config.DefaultResponseTableName, "Database table used to store response data")
+	flag.StringVar(&responseTables, "response-tables", "", "Comma separated list of database tables used to store response data")
 	// Oauth2 Client Credentials
 	flag.StringVar(&oauth2ClientID, "oauth2-client-id", "", "Oauth2 Client ID")
 	flag.StringVar(&oauth2ClientSecret, "oauth2-client-secret", "", "Oauth2 Client Secret")
@@ -57,39 +58,50 @@ func main() {
 	flag.StringVar(&caFile, "ca-file", "", "Path to the CA file")
 	flag.Parse()
 
-	if len(strings.TrimSpace(responseTableName)) == 0 {
-		log.Fatalf("response-table cannot be empty")
-	}
-
 	args := flag.Args()
 	if len(args) != 1 {
 		log.Fatalf("Usage: %s <flags> [DSN]\n\nExample:\n\t%s file:example.db?_journal=WAL&_sync=NORMAL&_timeout=5000&_txlock=immediate\n", os.Args[0], os.Args[0])
 	}
 	dsn := args[0]
 
-	db, err := sql.Open("sqlite3", dsn)
+	sqlDB, err := sql.Open("sqlite3", dsn)
 	if err != nil {
 		log.Fatalf("cannot connect to the database: %v", err)
 	}
-	defer db.Close()
+	defer sqlDB.Close()
+
+	var tableList []string
+	if responseTables == "" {
+		tableList, err = db.ResponseTables(sqlDB)
+		if err != nil {
+			log.Fatalf("discovery response tables: %v", err)
+		}
+	} else {
+		tableList = strings.Split(responseTables, ",")
 
-	_, err = db.Exec(fmt.Sprintf("CREATE VIRTUAL TABLE temp.http_refresh USING http_request(%s)", opts()))
-	if err != nil {
-		log.Fatalf("error creating virtual table: %v", err)
 	}
 
-	stmt, err := db.Prepare(fmt.Sprintf(`INSERT INTO temp.http_refresh(url) 
+	stmts := make(map[string]*sql.Stmt)
+	for _, responseTableName := range tableList {
+		_, err = sqlDB.Exec(fmt.Sprintf("CREATE VIRTUAL TABLE temp.%s_refresh USING http_request(%s)", responseTableName, opts(responseTableName)))
+		if err != nil {
+			log.Fatalf("error creating virtual table: %v", err)
+		}
+
+		stmt, err := sqlDB.Prepare(fmt.Sprintf(`INSERT INTO temp.%s_refresh(url) 
 	SELECT url FROM %s
-	WHERE url LIKE ? AND unixepoch() - unixepoch(timestamp) > ?`, responseTableName))
-	if err != nil {
-		log.Fatalf("error preparing statement: %v", err)
+	WHERE url LIKE ? AND unixepoch() - unixepoch(timestamp) > ?`, responseTableName, responseTableName))
+		if err != nil {
+			log.Fatalf("error preparing statement: %v", err)
+		}
+		defer stmt.Close()
+		stmts[responseTableName] = stmt
 	}
-	defer stmt.Close()
 
 	done := make(chan os.Signal, 1)
 	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
 
-	refreshData(stmt)
+	refreshData(stmts)
 	if interval == 0 {
 		return
 	}
@@ -102,27 +114,29 @@ func main() {
 	for {
 		select {
 		case <-ticker.C:
-			refreshData(stmt)
+			refreshData(stmts)
 		case <-done:
 			return
 		}
 	}
 }
 
-func refreshData(stmt *sql.Stmt) {
+func refreshData(stmts map[string]*sql.Stmt) {
 	slog.Info("starting data verification")
 
-	res, err := stmt.Exec(matchURL, ttl)
-	if err != nil {
-		slog.Error("error refreshing data", "error", err)
-		return
-	}
+	for tableName, stmt := range stmts {
+		res, err := stmt.Exec(matchURL, ttl)
+		if err != nil {
+			slog.Error("error refreshing data", "error", err, "table", tableName)
+			continue
+		}
 
-	rowsAffected, _ := res.RowsAffected()
-	slog.Info("verification finished", "rows_affected", rowsAffected)
+		rowsAffected, _ := res.RowsAffected()
+		slog.Info("verification finished", "rows_affected", rowsAffected, "table", tableName)
+	}
 }
 
-func opts() string {
+func opts(responseTableName string) string {
 	opts := make([]string, 0)
 	opts = append(opts, fmt.Sprintf("%s='%d'", config.Timeout, timeout))
 	opts = append(opts, fmt.Sprintf("%s='%v'", config.Insecure, insecure))