path: root/main.go

package main

import (
	"bytes"
	"crypto/sha1"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

var (
	db *sql.DB // sqlite database
	gd string  // gallery directory
)

func openDB(directory string) error {
	var err error
	db, err = sql.Open("sqlite3",
		"file:"+filepath.Join(directory, "gallery.db?_foreign_keys=1"))
	gd = directory
	return err
}

func imagePath(sha1 string) string {
	return filepath.Join(gd, "images", sha1[:2], sha1)
}

func thumbPath(sha1 string) string {
	return filepath.Join(gd, "thumbs", sha1[:2], sha1+".webp")
}

func dbCollect(query string) ([]string, error) {
	rows, err := db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var result []string
	for rows.Next() {
		var sha1 string
		if err := rows.Scan(&sha1); err != nil {
			return nil, err
		}
		result = append(result, sha1)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return result, nil
}

// cmdInit initializes a "gallery directory" that contains gallery.sqlite,
// images, thumbs.
func cmdInit(args []string) error {
	if len(args) != 1 {
		return errors.New("usage: GD")
	}

	if err := openDB(args[0]); err != nil {
		return err
	}
	if _, err := db.Exec(initializeSQL); err != nil {
		return err
	}

	// XXX: There's technically no reason to keep images as symlinks,
	// we might just keep absolute paths in the database as well.
	if err := os.MkdirAll(filepath.Join(gd, "images"), 0755); err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Join(gd, "thumbs"), 0755); err != nil {
		return err
	}
	return nil
}

// cmdRun runs a web UI against GD on ADDRESS.
func cmdRun(args []string) error {
	if len(args) != 2 {
		return errors.New("usage: GD ADDRESS")
	}
	if err := openDB(args[0]); err != nil {
		return err
	}

	address := args[1]

	http.Handle("/", http.FileServer(http.Dir("public")))
	// TODO: These subdirectories should be indirect
	// (skip the hash subpath, don't require the .webp suffix).
	http.Handle("/images",
		http.FileServer(http.Dir(filepath.Join(gd, "images"))))
	http.Handle("/thumbs",
		http.FileServer(http.Dir(filepath.Join(gd, "thumbs"))))

	host, port, err := net.SplitHostPort(address)
	if err != nil {
		log.Println(err)
	} else if host == "" {
		log.Println("http://" + net.JoinHostPort("localhost", port))
	} else {
		log.Println("http://" + address)
	}

	s := &http.Server{
		Addr:           address,
		ReadTimeout:    60 * time.Second,
		WriteTimeout:   60 * time.Second,
		MaxHeaderBytes: 32 << 10,
	}
	return s.ListenAndServe()
}

func isImage(path string) (bool, error) {
	cmd := exec.Command("xdg-mime", "query", "filetype", path)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return false, err
	}
	if err := cmd.Start(); err != nil {
		return false, err
	}
	out, err := io.ReadAll(stdout)
	if err != nil {
		return false, err
	}
	if err := cmd.Wait(); err != nil {
		return false, err
	}
	return bytes.HasPrefix(out, []byte("image/")), nil
}

func importFunc(path string, d fs.DirEntry, err error) error {
	if err != nil || d.IsDir() {
		return err
	}

	// The input may be a relative path, and we want to remember it as such,
	// but symlinks for the images must be absolute.
	absPath, err := filepath.Abs(path)
	if err != nil {
		return err
	}

	// Skip videos, which ImageMagick can process, but we don't want it to,
	// so that they're not converted 1:1 to WebP.
	pathIsImage, err := isImage(path)
	if err != nil {
		return err
	}
	if !pathIsImage {
		return nil
	}

	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()

	s, err := f.Stat()
	if err != nil {
		return err
	}

	hash := sha1.New()
	_, err = io.CopyBuffer(hash, f, make([]byte, 65536))
	if err != nil {
		return err
	}

	hexSHA1 := hex.EncodeToString(hash.Sum(nil))
	pathImage := imagePath(hexSHA1)
	imageDirname, _ := filepath.Split(pathImage)
	if err := os.MkdirAll(imageDirname, 0755); err != nil {
		return err
	}
	if err := os.Symlink(absPath, pathImage); err != nil &&
		!errors.Is(err, fs.ErrExist) {
		return err
	}

	// TODO: This should all run in a transaction.
	if _, err = db.Exec(`INSERT INTO image(sha1) VALUES (?)
		ON CONFLICT(sha1) DO NOTHING`, hexSHA1); err != nil {
		return err
	}

	dbDirname, dbBasename := filepath.Split(path)
	_, err = db.Exec(`INSERT INTO entry(
		path, basename, mtime, sha1
	) VALUES (?, ?, ?, ?)`, dbDirname, dbBasename, s.ModTime().Unix(), hexSHA1)
	return err
}

// cmdImport adds files to the "entry" table.
func cmdImport(args []string) error {
	if len(args) < 1 {
		return errors.New("usage: GD ROOT...")
	}
	if err := openDB(args[0]); err != nil {
		return err
	}

	// TODO: This would better be done in parallel (making hashes).
	// TODO: Show progress in some manner. Perhaps port my propeller code.
	for _, name := range args[1:] {
		if err := filepath.WalkDir(name, importFunc); err != nil {
			return err
		}
	}
	return nil
}

// cmdSync is like import, but clears the "entry" table beforehands.
func cmdSync(args []string) error {
	if len(args) < 1 {
		return errors.New("usage: GD ROOT...")
	}
	if err := openDB(args[0]); err != nil {
		return err
	}

	// TODO
	return nil
}

// cmdCheck checks if all files tracked in the DB are accessible.
func cmdCheck(args []string) error {
	if len(args) != 1 {
		return errors.New("usage: GD")
	}
	if err := openDB(args[0]); err != nil {
		return err
	}

	// TODO: Check if all hashes of DB entries have a statable image file,
	// and that all images with thumb{w,h} have a thumbnail file. Perhaps.
	return nil
}

func makeThumbnail(pathImage, pathThumb string) (int, int, error) {
	thumbDirname, _ := filepath.Split(pathThumb)
	if err := os.MkdirAll(thumbDirname, 0755); err != nil {
		return 0, 0, err
	}

	cmd := exec.Command("convert", pathImage, "-coalesce", "-colorspace", "RGB",
		"-auto-orient", "-strip", "-resize", "256x128>", "-colorspace", "sRGB",
		"-format", "%w %h", "+write", "info:", pathThumb)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return 0, 0, err
	}
	if err := cmd.Start(); err != nil {
		return 0, 0, err
	}
	out, err := io.ReadAll(stdout)
	if err != nil {
		return 0, 0, err
	}
	if err := cmd.Wait(); err != nil {
		return 0, 0, err
	}

	var w, h int
	_, err = fmt.Fscanf(bytes.NewReader(out), "%d %d", &w, &h)
	return w, h, err
}

// cmdThumbnail generates missing thumbnails, in parallel.
func cmdThumbnail(args []string) error {
	if len(args) < 1 {
		return errors.New("usage: GD [SHA1...]")
	}
	if err := openDB(args[0]); err != nil {
		return err
	}

	hexSHA1 := args[1:]
	if len(hexSHA1) == 0 {
		// Get all unique images in the database with no thumbnail.
		var err error
		hexSHA1, err = dbCollect(`SELECT sha1 FROM image
			WHERE thumbw IS NULL OR thumbh IS NULL`)
		if err != nil {
			return err
		}
	}

	// TODO: Try to run the thumbnailer in parallel, somehow.
	// TODO: Show progress in some manner. Perhaps port my propeller code.
	for _, sha1 := range hexSHA1 {
		pathImage := imagePath(sha1)
		pathThumb := thumbPath(sha1)
		w, h, err := makeThumbnail(pathImage, pathThumb)
		if err != nil {
			return err
		}

		_, err = db.Exec(`UPDATE image SET thumbw = ?, thumbh = ?
			WHERE sha1 = ?`, w, h, sha1)
		if err != nil {
			return err
		}
	}
	return nil
}

func makeDhash(hasher, pathThumb string) (uint64, error) {
	cmd := exec.Command(hasher, pathThumb)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return 0, err
	}
	if err := cmd.Start(); err != nil {
		return 0, err
	}
	out, err := io.ReadAll(stdout)
	if err != nil {
		return 0, err
	}
	if err := cmd.Wait(); err != nil {
		return 0, err
	}

	var hash uint64
	_, err = fmt.Fscanf(bytes.NewReader(out), "%x", &hash)
	return hash, err
}

// cmdDhash generates perceptual hash from thumbnails.
func cmdDhash(args []string) error {
	if len(args) < 1 {
		return errors.New("usage: GD HASHER [SHA1...]")
	}
	if err := openDB(args[0]); err != nil {
		return err
	}

	hasher, hexSHA1 := args[1], args[2:]
	if len(hexSHA1) == 0 {
		var err error
		hexSHA1, err = dbCollect(`SELECT sha1 FROM image WHERE dhash IS NULL`)
		if err != nil {
			return err
		}
	}

	// TODO: Try to run the hasher in parallel, somehow.
	// TODO: Show progress in some manner. Perhaps port my propeller code.
	for _, sha1 := range hexSHA1 {
		pathThumb := thumbPath(sha1)
		hash, err := makeDhash(hasher, pathThumb)
		if err != nil {
			return err
		}

		_, err = db.Exec(`UPDATE image SET dhash = ? WHERE sha1 = ?`,
			int64(hash), sha1)
		if err != nil {
			return err
		}
	}
	return nil
}

var commands = map[string]struct {
	handler func(args []string) error
}{
	"init":      {cmdInit},
	"run":       {cmdRun},
	"import":    {cmdImport},
	"sync":      {cmdSync},
	"check":     {cmdCheck},
	"thumbnail": {cmdThumbnail},
	"dhash":     {cmdDhash},
}

func main() {
	if len(os.Args) <= 2 {
		log.Fatalln("Missing arguments")
	}

	cmd, ok := commands[os.Args[1]]
	if !ok {
		log.Fatalln("Unknown command: " + os.Args[1])
	}

	err := cmd.handler(os.Args[2:])

	// Note that the database object has a closing finalizer,
	// we just additionally print any errors coming from there.
	if db != nil {
		if err := db.Close(); err != nil {
			log.Println(err)
		}
	}

	if err != nil {
		log.Fatalln(err)
	}
}