Use firewalls (#5)

* use firewalls

Signed-off-by: Flipez <[email protected]>

* add delete and clear

Signed-off-by: Flipez <[email protected]>

* update readme

Signed-off-by: Flipez <[email protected]>

---------

Signed-off-by: Flipez <[email protected]>

Author: Flipez

README.md

@@ -1,6 +1,8 @@
-# key/value store based on Hetzner Cloud Labels
-Implements a simple key/value store by utilizing the ssh-key resource and its labels.
-As per API spec it supports keys and values with a max lenght of 63 chars.
+# Slow key/value store based on Hetzner Cloud Metadata
+Implements a simple key/value store by utilizing the firewall resource and its rule descriptions.
+As per API spec it supports up to 500 rules with 255 bytes description each.
+The data is encoded using ![gob](https://pkg.go.dev/encoding/gob) and compressed using ![zstd](https://pkg.go.dev/github.com/klauspost/compress/zstd).
+The actual number of keys and size of values therefore depends on how well they can be compressed.
 hetzner-kv supports multiple "databases" provided by a flag.
 
 It uses the token specified via env variable `HCLOUD_TOKEN`
@@ -15,13 +17,16 @@ USAGE:
    hcloud-kv [global options] command [command options] [arguments...]
 
 COMMANDS:
-   init, i  initializes a new database
-   set, s   sets a key
-   get, g   get a value from given key
-   list, l  list all keys
-   help, h  Shows a list of commands or help for one command
+   init, i    initializes a new database
+   set, s     sets a key
+   get, g     get a value from given key
+   list, l    list all keys
+   delete, d  delete a given key
+   clear, c   delete all keys
+   help, h    Shows a list of commands or help for one command
 
 GLOBAL OPTIONS:
    --db value  database to use (default: "0")
+   --no-info   Do not print db usage information (default: false)
    --help, -h  show help
 ```

database.go

@@ -1,68 +1,89 @@
 package main
 
 import (
+	"bytes"
 	"context"
-	"encoding/json"
+	"encoding/base64"
+	"encoding/gob"
+	"fmt"
 	"log"
+	"net"
+	"strings"
 
 	"github.com/hetznercloud/hcloud-go/hcloud"
+	"github.com/klauspost/compress/zstd"
 )
 
 type Database struct {
-	Store   map[string]string
-	Name    string
-	Client  *hcloud.Client
-	Context context.Context
-	Self    *hcloud.SSHKey
-	NoInfo  bool
+	Store           map[string]string
+	Name            string
+	Client          *hcloud.Client
+	Context         context.Context
+	Self            *hcloud.Firewall
+	NoInfo          bool
+	LastEncodedSize int
 }
 
 func (d *Database) Init() {
-	keyPlaceholder := "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIAUqWKFtn3P3p0tOWXMgkfc7aTc5Z17+LSlf50X/ep/Z"
+	_, response, err := d.Client.Firewall.Create(d.Context, hcloud.FirewallCreateOpts{Name: d.Name})
 
-	database, response, err := d.Client.SSHKey.Create(d.Context, hcloud.SSHKeyCreateOpts{Name: d.Name, PublicKey: keyPlaceholder})
-	if err != nil && response.StatusCode == 409 {
-		log.Fatalf("database %s already exists", d.Name)
-	} else if err != nil {
-		log.Fatalf("unhandled error: %s\n", err)
+	if err != nil {
+		if response != nil && response.StatusCode == 409 {
+			log.Printf("database %s already exists", d.Name)
+		} else {
+			log.Fatalf("unhandled error: %s\n", err)
+		}
+	} else {
+		log.Printf("created new database: %s", d.Name)
 	}
 
-	d.Store = database.Labels
+	d.Fetch()
 }
 
-func (d *Database) Fetch() *hcloud.SSHKey {
-	db, _, err := d.Client.SSHKey.Get(d.Context, d.Name)
-	d.Self = db
-
+func (d *Database) Fetch() *hcloud.Firewall {
+	db, _, err := d.Client.Firewall.Get(d.Context, d.Name)
 	if err != nil {
 		log.Fatalf("error retrieving database: %s\n", err)
 	}
+	d.Self = db
 
 	if d.Self != nil {
-		d.Store = d.Self.Labels
-	} else {
-		log.Fatalf("database %s not found", d.Name)
+		store, size, err := rulesToMap(d.Self.Rules)
+		if err != nil {
+			log.Printf("could not parse rules (db might be empty or old format): %s", err)
+			if d.Store == nil {
+				d.Store = make(map[string]string)
+			}
+		} else {
+			d.Store = store
+			d.LastEncodedSize = size
+		}
 	}
 
-	checkSize(d)
-
+	checkSize(d, d.LastEncodedSize)
 	return d.Self
 }
 
 func (d *Database) Set(key, value string) bool {
 	d.Fetch()
+
+	if d.Store == nil {
+		d.Store = make(map[string]string)
+	}
+
 	d.Store[key] = value
 
-	database, _, err := d.Client.SSHKey.Update(d.Context, d.Self, hcloud.SSHKeyUpdateOpts{Labels: d.Store})
+	firewallRules, encodedLength, err := mapToRules(d.Store)
+
+	_, _, err = d.Client.Firewall.SetRules(d.Context, d.Self, hcloud.FirewallSetRulesOpts{Rules: firewallRules})
 
 	if err != nil {
 		log.Fatalf("error updating db: %s\n", err)
 	}
 
-	d.Store = database.Labels
-
+	d.LastEncodedSize = encodedLength
 	log.Println("OK")
-	checkSize(d)
+	checkSize(d, encodedLength)
 
 	return true
 }
@@ -75,21 +96,146 @@ func (d *Database) Get(key string) string {
 func (d *Database) List() []string {
 	d.Fetch()
 
-	keys := make([]string, len(d.Store))
+	keys := make([]string, 0, len(d.Store))
 
-	i := 0
 	for k := range d.Store {
-		keys[i] = k
-		i++
+		keys = append(keys, k)
 	}
 
 	return keys
 }
 
-func checkSize(db *Database) {
-	if !db.NoInfo {
-		jsonStr, _ := json.Marshal(db.Store)
-		log.Printf("[Info] Database usage: %.2f%%", float64(len(jsonStr))/float64(maxDBBytes)*100)
+func (d *Database) Delete(key string) bool {
+	d.Fetch()
+
+	if _, exists := d.Store[key]; !exists {
+		log.Printf("Key '%s' not found, nothing to delete.", key)
+		return false
 	}
+
+	delete(d.Store, key)
+
+	firewallRules, encodedLength, err := mapToRules(d.Store)
+	if err != nil {
+		log.Fatalf("Error encoding data after delete: %s", err)
+	}
+
+	_, _, err = d.Client.Firewall.SetRules(d.Context, d.Self, hcloud.FirewallSetRulesOpts{
+		Rules: firewallRules,
+	})
+
+	if err != nil {
+		log.Fatalf("error updating db during delete: %s\n", err)
+	}
+
+	d.LastEncodedSize = encodedLength
+	log.Println("Deleted key:", key)
+	checkSize(d, encodedLength)
+
+	return true
 }
 
+func (d *Database) Clear() {
+	d.Fetch()
+
+	if d.Self == nil {
+		log.Fatal("Could not clear database: Firewall not found or not initialized.")
+	}
+
+	d.Store = make(map[string]string)
+
+	_, _, err := d.Client.Firewall.SetRules(d.Context, d.Self, hcloud.FirewallSetRulesOpts{
+		Rules: []hcloud.FirewallRule{},
+	})
+
+	if err != nil {
+		log.Fatalf("Failed to clear database: %s", err)
+	}
+
+	d.LastEncodedSize = 0
+	log.Println("Database cleared successfully.")
+}
+
+func checkSize(db *Database, currentLength int) {
+	if db.NoInfo {
+		return
+	}
+
+	maxChars := 500 * 255
+	usage := (float64(currentLength) / float64(maxChars)) * 100
+
+	log.Printf("[Info] Storage: %d/%d chars (%.2f%% used)", currentLength, maxChars, usage)
+}
+
+func rulesToMap(rules []hcloud.FirewallRule) (map[string]string, int, error) {
+	store := map[string]string{}
+	var sb strings.Builder
+	sb.Grow(len(rules) * 255)
+
+	for _, rule := range rules {
+		if rule.Description != nil {
+			sb.WriteString(*rule.Description)
+		}
+	}
+	tempString := sb.String()
+	totalLength := len(tempString)
+	if totalLength == 0 {
+		return store, 0, nil
+	}
+
+	decodedBytes, err := base64.StdEncoding.DecodeString(tempString)
+	if err != nil {
+		return store, totalLength, err
+	}
+
+	zr, err := zstd.NewReader(bytes.NewReader(decodedBytes))
+	if err != nil {
+		return store, totalLength, err
+	}
+	defer zr.Close()
+
+	err = gob.NewDecoder(zr).Decode(&store)
+	return store, totalLength, err
+}
+
+func mapToRules(store map[string]string) ([]hcloud.FirewallRule, int, error) {
+	var gobBuf bytes.Buffer
+	if err := gob.NewEncoder(&gobBuf).Encode(store); err != nil {
+		return nil, 0, err
+	}
+
+	var compressedBuf bytes.Buffer
+	zw, _ := zstd.NewWriter(&compressedBuf)
+	zw.Write(gobBuf.Bytes())
+	zw.Close()
+
+	encodedString := base64.StdEncoding.EncodeToString(compressedBuf.Bytes())
+	totalLength := len(encodedString)
+
+	var rules []hcloud.FirewallRule
+	limit := 255
+	_, dummyNet, _ := net.ParseCIDR("0.0.0.0/32")
+
+	for i := 0; i < len(encodedString); i += limit {
+		end := i + limit
+		if end > len(encodedString) {
+			end = len(encodedString)
+		}
+
+		chunk := encodedString[i:end] // Direct slice is safer for ASCII Base64
+
+		rules = append(rules, hcloud.FirewallRule{
+			Description: hcloud.Ptr(chunk),
+			Direction:   hcloud.FirewallRuleDirectionIn,
+			Protocol:    hcloud.FirewallRuleProtocolTCP,
+			Port:        hcloud.Ptr("80"),
+			SourceIPs:   []net.IPNet{*dummyNet},
+		})
+	}
+
+	for i, r := range rules {
+		fmt.Printf("Rule %d: Desc Length: %d, Protocol: %s\n", i, len(*r.Description), r.Protocol)
+	}
+
+	return rules, totalLength, nil
+}

go.mod

@@ -4,14 +4,14 @@ go 1.24.0
 
 require (
 	github.com/hetznercloud/hcloud-go v1.59.2
+	github.com/klauspost/compress v1.18.2
 	github.com/urfave/cli/v2 v2.27.7
 )
 
 require (
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.7 // indirect
-	github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/prometheus/client_golang v1.23.2 // indirect
 	github.com/prometheus/client_model v0.6.2 // indirect