db.go

package main

import (
	"bytes"
	"database/sql"
	"encoding/gob"
	"errors"
	"strconv"
	"time"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"
	_ "github.com/mattn/go-sqlite3"
)

var ErrNotAuthorized = errors.New("Not authorized")

// ZerodropEntry is a page entry.
type ZerodropEntry struct {
	Name                 string    // The request URI used to access this entry
	URL                  string    // The URL that this entry references
	Filename             string    // The location of the file in the uploads directory
	ContentType          string    // The MIME type to serve as Content-Type header
	Redirect             bool      // Indicates whether to redirect instead of proxy
	Creation             time.Time // The time this entry was created
	AccessRedirectOnDeny string    // Entry to redirect to if entry is blacklisted or expired
	AccessBlacklist      Blacklist // Blacklist
	AccessBlacklistCount int       // Number of requests that have been caught by the blacklist
	AccessExpire         bool      // Indicates whether to expire after finite access
	AccessExpireCount    int       // The number of requests on this entry before expiry
	AccessCount          int       // The number of times this has been accessed
	AccessTrain          bool      // Whether training is active
}

// ZerodropDB represents a database connection.
// TODO: Use a persistent backend.
type ZerodropDB struct {
	*sql.DB

	// Accessors
	GetStmt       *sql.Stmt
	ListStmt      *sql.Stmt
	ListTokenStmt *sql.Stmt

	// Restricted Mutators
	UpdateCheckTokenStmt *sql.Stmt
	DeleteStmt           *sql.Stmt
	ClearStmt            *sql.Stmt

	// Admin Mutators
	AdminUpdateStmt *sql.Stmt
	AdminDeleteStmt *sql.Stmt
	AdminClearStmt  *sql.Stmt
}

// Connect opens a connection to the backend.
func (d *ZerodropDB) Connect(driver, source string) error {
	db, err := sql.Open(driver, source)
	if err != nil {
		return err
	}

	gob.Register(&ZerodropEntry{})

	if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS entries (
		name TEXT PRIMARY KEY NOT NULL,
		token TEXT NOT NULL,
		creation INTEGER NOT NULL,
		gob BLOB NOT NULL
	)`); err != nil {
		return err
	}

	// Accessors
	d.GetStmt, err = db.Prepare(
		`SELECT gob FROM entries WHERE name = ?`)
	if err != nil {
		return err
	}
	d.ListStmt, err = db.Prepare(
		`SELECT gob FROM entries ORDER BY creation DESC`)
	if err != nil {
		return err
	}
	d.ListTokenStmt, err = db.Prepare(
		`SELECT gob FROM entries WHERE token = ? ORDER BY creation DESC`)
	if err != nil {
		return err
	}

	// Restricted Mutators
	d.UpdateCheckTokenStmt, err = db.Prepare(
		`SELECT token FROM entries WHERE name = ?`)
	if err != nil {
		return err
	}
	d.DeleteStmt, err = db.Prepare(
		`DELETE FROM entries WHERE name = ? AND token = ?`)
	if err != nil {
		return err
	}
	d.ClearStmt, err = db.Prepare(
		`DELETE FROM entries WHERE token = ?`)
	if err != nil {
		return err
	}

	// Admin Mutators
	d.AdminUpdateStmt, err = db.Prepare(
		`REPLACE INTO entries (name, token, creation, gob) VALUES (?, ?, ?, ?)`)
	if err != nil {
		return err
	}
	d.AdminDeleteStmt, err = db.Prepare(
		`DELETE FROM entries WHERE name = ?`)
	if err != nil {
		return err
	}
	d.AdminClearStmt, err = db.Prepare(
		`DELETE FROM entries`)
	if err != nil {
		return err
	}

	d.DB = db
	return nil
}

// Get returns the entry with the specified name.
func (d *ZerodropDB) Get(name string) (*ZerodropEntry, error) {
	var data []byte
	if err := d.GetStmt.QueryRow(name).Scan(&data); err != nil {
		return nil, err
	}

	var entry *ZerodropEntry
	dec := gob.NewDecoder(bytes.NewReader(data))
	if err := dec.Decode(&entry); err != nil {
		return nil, err
	}

	return entry, nil
}

// List returns a slice of all entries sorted by creation time,
// with the most recent first.
func (d *ZerodropDB) List(token string) ([]*ZerodropEntry, error) {
	list := []*ZerodropEntry{}

	var err error
	var rows *sql.Rows

	if token == "" {
		rows, err = d.ListStmt.Query()
	} else {
		rows, err = d.ListTokenStmt.Query(token)
	}
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var data []byte
		if err := rows.Scan(&data); err != nil {
			return nil, err
		}

		var entry *ZerodropEntry
		dec := gob.NewDecoder(bytes.NewReader(data))
		if err := dec.Decode(&entry); err != nil {
			return nil, err
		}

		list = append(list, entry)
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	return list, nil
}

// Update adds an entry to the database.
func (d *ZerodropDB) Update(entry *ZerodropEntry, claims *AdminClaims) error {
	var token string

	err := d.UpdateCheckTokenStmt.QueryRow(entry.Name).Scan(&token)
	if err != nil {
		// The entry does not exist.
		token = claims.Token
	} else if !claims.Admin && token != claims.Token {
		// The entry exists and the tokens do not match.
		return ErrNotAuthorized
	}

	var buffer bytes.Buffer
	enc := gob.NewEncoder(&buffer)
	if err := enc.Encode(entry); err != nil {
		return err
	}

	if _, err := d.AdminUpdateStmt.Exec(entry.Name, token, entry.Creation.Unix(), buffer.Bytes()); err != nil {
		return err
	}

	return nil
}

// Remove removes an entry from the database with the specified token.
func (d *ZerodropDB) Remove(name string, claims *AdminClaims) error {
	if claims.Admin {
		if _, err := d.AdminDeleteStmt.Exec(name); err != nil {
			return err
		}
		return nil
	}

	if _, err := d.DeleteStmt.Exec(name, claims.Token); err != nil {
		return err
	}

	return nil
}

// Clear resets the database by removing all entries with the specified token.
func (d *ZerodropDB) Clear(claims *AdminClaims) error {
	if claims.Admin {
		if _, err := d.AdminClearStmt.Exec(); err != nil {
			return err
		}

		return nil
	}

	if _, err := d.ClearStmt.Exec(claims.Token); err != nil {
		return err
	}

	return nil
}

// IsExpired returns true if the entry is expired
func (e *ZerodropEntry) IsExpired() bool {
	return e.AccessExpire && (e.AccessCount >= e.AccessExpireCount)
}

// SetTraining sets the AccessTrain flag
func (e *ZerodropEntry) SetTraining(train bool) {
	e.AccessTrain = train
}

// Access increases the access count for an entry.
func (e *ZerodropEntry) Access() {
	e.AccessCount++
}

func (e *ZerodropEntry) String() string {
	return strconv.Quote(e.Name)
}