styx/proxy/stats/stats.go

package stats

import (
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/user"
	"path/filepath"
	"time"

	"git.maze.io/maze/styx/internal/log"
	_ "github.com/mattn/go-sqlite3"
)

type Stats struct {
	db *sql.DB
}

func New() (*Stats, error) {
	u, err := user.Current()
	if err != nil {
		return nil, err
	}

	path := filepath.Join(u.HomeDir, ".styx", "stats.db")
	if err = os.MkdirAll(filepath.Dir(path), 0o750); err != nil {
		return nil, err
	}

	db, err := sql.Open("sqlite3", path+"?_journal_mode=WAL")
	if err != nil {
		return nil, err
	}

	for _, table := range []string{
		createLog,
		createDomainStat,
		createStatusStat,
	} {
		if _, err = db.Exec(table); err != nil {
			return nil, err
		}
	}

	return &Stats{db: db}, nil
}

func (s *Stats) AddLog(entry *Log) error {
	var (
		request  []byte
		response []byte
		err      error
	)
	if request, err = json.Marshal(entry.Request); err != nil {
		return err
	}
	if response, err = json.Marshal(entry.Response); err != nil {
		return err
	}

	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	stmt, err := tx.Prepare("insert into styx_log(client_ip, request, response) values(?, ?, ?)")
	if err != nil {
		return err
	}
	defer stmt.Close()
	if _, err = stmt.Exec(entry.ClientIP, request, response); err != nil {
		return err
	}
	return tx.Commit()
}

func (s *Stats) QueryLog(offset, limit int) ([]*Log, error) {
	if limit == 0 {
		limit = 50
	}

	rows, err := s.db.Query("select dt, client_ip, request, response from styx_log limit ?, ?", offset, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var logs []*Log
	for rows.Next() {
		var entry = new(Log)
		if err = rows.Scan(&entry.Time, &entry.ClientIP, &entry.Request, &entry.Response); err != nil {
			return nil, err
		}
		logs = append(logs, entry)
	}

	return logs, nil
}

type Status struct {
	Code  int `json:"code"`
	Count int `json:"count"`
}

var timeZero time.Time

func (s *Stats) QueryStatus(since time.Time) ([]*Status, error) {
	if since.Equal(timeZero) {
		since = time.Now().Add(-24 * time.Hour)
	}

	rows, err := s.db.Query("select response->'status', count(*) from styx_log where dt >= ? group by response->'status' order by response->'status'", since)
	if err != nil {
		return nil, err
	}

	var stats []*Status
	for rows.Next() {
		var entry = new(Status)
		if err = rows.Scan(&entry.Code, &entry.Count); err != nil {
			return nil, err
		}
		stats = append(stats, entry)
	}
	return stats, nil
}

const createLog = `CREATE TABLE IF NOT EXISTS styx_log (
	id        INT PRIMARY KEY,
	dt        DATETIME DEFAULT CURRENT_TIMESTAMP,
	client_ip TEXT NOT NULL,
	request   JSONB NOT NULL,
	response  JSONB NOT NULL
);`

type Log struct {
	Time     time.Time `json:"time"`
	ClientIP string    `json:"client_ip"`
	Request  *Request  `json:"request"`
	Response *Response `json:"response"`
}

type Request struct {
	URL    string      `json:"url"`
	Host   string      `json:"host"`
	Method string      `json:"method"`
	Proto  string      `json:"proto"`
	Header http.Header `json:"header"`
}

func (r *Request) Scan(value any) error {
	switch v := value.(type) {
	case string:
		return json.Unmarshal([]byte(v), r)
	case []byte:
		return json.Unmarshal(v, r)
	default:
		log.Error().Str("type", fmt.Sprintf("%T", value)).Msg("scan request unknown type")
		return nil
	}
}

func (r *Request) Value() (driver.Value, error) {
	b, err := json.Marshal(r)
	return string(b), err
}

func FromRequest(r *http.Request) *Request {
	return &Request{
		URL:    r.URL.String(),
		Host:   r.Host,
		Method: r.Method,
		Proto:  r.Proto,
		Header: r.Header,
	}
}

type Response struct {
	Status int         `json:"status"`
	Size   int64       `json:"size"`
	Header http.Header `json:"header"`
}

func (r *Response) Scan(value any) error {
	switch v := value.(type) {
	case string:
		return json.Unmarshal([]byte(v), r)
	case []byte:
		return json.Unmarshal(v, r)
	default:
		log.Error().Str("type", fmt.Sprintf("%T", value)).Msg("scan response unknown type")
		return nil
	}
}

func (r *Response) Value() (driver.Value, error) {
	b, err := json.Marshal(r)
	return string(b), err
}

func (r *Response) SetSize(size int64) *Response {
	r.Size = size
	return r
}

func FromResponse(r *http.Response) *Response {
	return &Response{
		Status: r.StatusCode,
		Header: r.Header,
	}
}

const createStatusStat = `CREATE TABLE IF NOT EXISTS styx_stat_status (
	id     INT PRIMARY KEY,
	dt     DATETIME DEFAULT CURRENT_TIMESTAMP,
	status INT NOT NULL
);`

const createDomainStat = `CREATE TABLE IF NOT EXISTS styx_stat_domain (
	id     INT PRIMARY KEY,
	dt     DATETIME DEFAULT CURRENT_TIMESTAMP,
	domain TEXT NOT NULL
);`