Initial commit

1 files changed, 105 insertions, 0 deletions

diff --git a/database/connections.go b/database/connections.go
new file mode 100644
index 0000000..e16a17e
--- /dev/null
+++ b/database/connections.go
@@ -0,0 +1,105 @@
+package database
+
+import (
+	"database/sql"
+	"errors"
+	"time"
+
+	_ "github.com/mattn/go-sqlite3"
+)
+
+// connectionFreshness is how long a device counts as present after we last saw
+// it in the ARP table. It smooths over the kernel briefly aging an entry out.
+const connectionFreshness = 60 * time.Second
+
+type Connection struct {
+	IP        string
+	Mac       string
+	UpdatedAt time.Time
+}
+
+// RefreshConnections upserts the current ARP snapshot, bumping updated_at for
+// devices still present, and prunes ones not seen within the freshness window.
+func RefreshConnections(db *sql.DB, ipToMac map[string]string) error {
+	tx, err := db.Begin()
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+
+	stmt, err := tx.Prepare(`INSERT INTO connections (ip, mac, updated_at) VALUES (?, ?, ?)
+		ON CONFLICT(ip) DO UPDATE SET mac = excluded.mac, updated_at = excluded.updated_at`)
+	if err != nil {
+		return err
+	}
+	defer stmt.Close()
+
+	now := time.Now()
+	for ip, mac := range ipToMac {
+		if _, err := stmt.Exec(ip, mac, now); err != nil {
+			return err
+		}
+	}
+
+	if _, err := tx.Exec("DELETE FROM connections WHERE updated_at < ?", now.Add(-connectionFreshness)); err != nil {
+		return err
+	}
+
+	return tx.Commit()
+}
+
+func MacForIP(db *sql.DB, ip string) (string, error) {
+	var mac string
+	err := db.QueryRow(
+		"SELECT mac FROM connections WHERE ip = ? AND updated_at > ?",
+		ip, time.Now().Add(-connectionFreshness),
+	).Scan(&mac)
+	if errors.Is(err, sql.ErrNoRows) {
+		return "", nil
+	}
+	return mac, err
+}
+
+// NetworkUserIDs returns the set of users whose MAC was seen on the LAN within
+// the freshness window, whether or not they're on the chat.
+func NetworkUserIDs(db *sql.DB) (map[string]bool, error) {
+	rows, err := db.Query(
+		"SELECT DISTINCT m.user_id FROM macs m JOIN connections c ON c.mac = m.mac WHERE c.updated_at > ?",
+		time.Now().Add(-connectionFreshness),
+	)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	ids := map[string]bool{}
+	for rows.Next() {
+		var id string
+		if err := rows.Scan(&id); err != nil {
+			return nil, err
+		}
+		ids[id] = true
+	}
+	return ids, rows.Err()
+}
+
+func ListAliveConnections(db *sql.DB) ([]Connection, error) {
+	rows, err := db.Query(
+		"SELECT ip, mac, updated_at FROM connections WHERE updated_at > ? ORDER BY ip",
+		time.Now().Add(-connectionFreshness),
+	)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	conns := []Connection{}
+	for rows.Next() {
+		var c Connection
+		if err := rows.Scan(&c.IP, &c.Mac, &c.UpdatedAt); err != nil {
+			return nil, err
+		}
+		conns = append(conns, c)
+	}
+	return conns, rows.Err()
+}