tomswallapi/internal/store/longtail.go

package store

import (
	"context"
	"errors"

	"github.com/jackc/pgx/v5"

	"git.unkin.net/unkin/tomswallapi/internal/model"
)

// ---- Policies (table exists from 0001; ListPolicies is in store.go) ---------

func (s *Store) GetPolicy(ctx context.Context, id int64) (model.Policy, error) {
	var p model.Policy
	err := s.pool.QueryRow(ctx,
		`SELECT id, priority, source, dest, action, log FROM policies WHERE id = $1`, id,
	).Scan(&p.ID, &p.Priority, &p.Source, &p.Dest, &p.Action, &p.Log)
	if errors.Is(err, pgx.ErrNoRows) {
		return p, ErrNotFound
	}
	return p, err
}

func (s *Store) CreatePolicy(ctx context.Context, p model.Policy) (int64, error) {
	var id int64
	err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
		if err := tx.QueryRow(ctx, `
			INSERT INTO policies (priority, source, dest, action, log)
			VALUES ($1, $2, $3, $4, $5) RETURNING id`,
			p.Priority, p.Source, p.Dest, p.Action, p.Log,
		).Scan(&id); err != nil {
			return err
		}
		return bump(ctx, tx)
	})
	return id, err
}

func (s *Store) DeletePolicy(ctx context.Context, id int64) error {
	return s.deleteOne(ctx, `DELETE FROM policies WHERE id = $1`, id)
}

// ---- Blrules ---------------------------------------------------------------

func (s *Store) ListBlrules(ctx context.Context) ([]model.BlruleRule, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT id, priority, action, source, dest, proto, dport, sport, log, comment
		FROM blrules ORDER BY priority, id`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []model.BlruleRule
	for rows.Next() {
		r, err := scanBlrule(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

func (s *Store) GetBlrule(ctx context.Context, id int64) (model.BlruleRule, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT id, priority, action, source, dest, proto, dport, sport, log, comment
		FROM blrules WHERE id = $1`, id)
	if err != nil {
		return model.BlruleRule{}, err
	}
	defer rows.Close()
	if !rows.Next() {
		return model.BlruleRule{}, ErrNotFound
	}
	return scanBlrule(rows)
}

func scanBlrule(rows pgx.Rows) (model.BlruleRule, error) {
	var r model.BlruleRule
	var dport, sport []byte
	if err := rows.Scan(&r.ID, &r.Priority, &r.Action, &r.Source, &r.Dest, &r.Proto, &dport, &sport, &r.Log, &r.Comment); err != nil {
		return r, err
	}
	if err := unmarshalStrings(dport, &r.DPort); err != nil {
		return r, err
	}
	return r, unmarshalStrings(sport, &r.SPort)
}

func (s *Store) CreateBlrule(ctx context.Context, r model.BlruleRule) (int64, error) {
	dport, _ := jsonb(r.DPort)
	sport, _ := jsonb(r.SPort)
	var id int64
	err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
		if err := tx.QueryRow(ctx, `
			INSERT INTO blrules (priority, action, source, dest, proto, dport, sport, log, comment)
			VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) RETURNING id`,
			r.Priority, r.Action, r.Source, r.Dest, r.Proto, dport, sport, r.Log, r.Comment,
		).Scan(&id); err != nil {
			return err
		}
		return bump(ctx, tx)
	})
	return id, err
}

func (s *Store) DeleteBlrule(ctx context.Context, id int64) error {
	return s.deleteOne(ctx, `DELETE FROM blrules WHERE id = $1`, id)
}

// ---- Conntrack -------------------------------------------------------------

func (s *Store) ListConntrack(ctx context.Context) ([]model.ConntrackRule, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT id, priority, action, source, dest, proto, dport, sport, chain, helper, "user", comment
		FROM conntrack ORDER BY priority, id`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []model.ConntrackRule
	for rows.Next() {
		r, err := scanConntrack(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

func (s *Store) GetConntrack(ctx context.Context, id int64) (model.ConntrackRule, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT id, priority, action, source, dest, proto, dport, sport, chain, helper, "user", comment
		FROM conntrack WHERE id = $1`, id)
	if err != nil {
		return model.ConntrackRule{}, err
	}
	defer rows.Close()
	if !rows.Next() {
		return model.ConntrackRule{}, ErrNotFound
	}
	return scanConntrack(rows)
}

func scanConntrack(rows pgx.Rows) (model.ConntrackRule, error) {
	var r model.ConntrackRule
	var dport, sport []byte
	if err := rows.Scan(&r.ID, &r.Priority, &r.Action, &r.Source, &r.Dest, &r.Proto, &dport, &sport, &r.Chain, &r.Helper, &r.User, &r.Comment); err != nil {
		return r, err
	}
	if err := unmarshalStrings(dport, &r.DPort); err != nil {
		return r, err
	}
	return r, unmarshalStrings(sport, &r.SPort)
}

func (s *Store) CreateConntrack(ctx context.Context, r model.ConntrackRule) (int64, error) {
	dport, _ := jsonb(r.DPort)
	sport, _ := jsonb(r.SPort)
	var id int64
	err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
		if err := tx.QueryRow(ctx, `
			INSERT INTO conntrack (priority, action, source, dest, proto, dport, sport, chain, helper, "user", comment)
			VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) RETURNING id`,
			r.Priority, r.Action, r.Source, r.Dest, r.Proto, dport, sport, r.Chain, r.Helper, r.User, r.Comment,
		).Scan(&id); err != nil {
			return err
		}
		return bump(ctx, tx)
	})
	return id, err
}

func (s *Store) DeleteConntrack(ctx context.Context, id int64) error {
	return s.deleteOne(ctx, `DELETE FROM conntrack WHERE id = $1`, id)
}