From 2b27cc57e34f17ceded04a76f4af89e0e8c4f66c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mika=C3=ABl=20Cluseau?= <mikael.cluseau@gmail.com>
Date: Sun, 27 Jul 2025 12:40:01 +0200
Subject: [PATCH] render context: add asset_download_token

---
 cmd/dkl-local-server/render-context.go  | 27 +++++++++++++
 cmd/dkl-local-server/ws-download-set.go | 50 ++++++++++++-------------
 2 files changed, 51 insertions(+), 26 deletions(-)

diff --git a/cmd/dkl-local-server/render-context.go b/cmd/dkl-local-server/render-context.go
index b70a7c0..8f57fe7 100644
--- a/cmd/dkl-local-server/render-context.go
+++ b/cmd/dkl-local-server/render-context.go
@@ -14,6 +14,7 @@ import (
 	"path"
 	"path/filepath"
 	"text/template"
+	"time"
 
 	cfsslconfig "github.com/cloudflare/cfssl/config"
 	restful "github.com/emicklei/go-restful"
@@ -237,6 +238,32 @@ func (ctx *renderContext) TemplateFuncs() map[string]any {
 				}
 			}
 
+			return
+		},
+		"asset_download_token": func(asset string, params ...string) (token string, err error) {
+			now := time.Now()
+			exp := now.Add(24 * time.Hour) // expire in 24h by default
+			if len(params) != 0 {
+				exp, err = parseCertDuration(params[0], now)
+				if err != nil {
+					return
+				}
+			}
+
+			set := DownloadSet{
+				Expiry: exp,
+				Items: []DownloadSetItem{
+					{
+						Kind:   "host",
+						Name:   ctx.Host.Name,
+						Assets: []string{asset},
+					},
+				},
+			}
+
+			privKey, _ := dlsSigningKeys()
+			token = set.Signed(privKey)
+
 			return
 		},
 	} {
diff --git a/cmd/dkl-local-server/ws-download-set.go b/cmd/dkl-local-server/ws-download-set.go
index 6a2d02b..e6c1b50 100644
--- a/cmd/dkl-local-server/ws-download-set.go
+++ b/cmd/dkl-local-server/ws-download-set.go
@@ -70,6 +70,28 @@ func (s *DownloadSet) Decode(encoded string) (err error) {
 	return
 }
 
+func (s DownloadSet) Signed(privKey ed25519.PrivateKey) string {
+	buf := new(bytes.Buffer)
+	{
+		setBytes := []byte(s.Encode())
+
+		w := lz4.NewWriter(buf)
+		w.Write(setBytes)
+		w.Close()
+	}
+
+	setBytes := buf.Bytes()
+	sig := ed25519.Sign(privKey, setBytes)
+
+	buf = bytes.NewBuffer(make([]byte, 0, 1+len(sig)+len(setBytes)))
+	buf.WriteByte(byte(len(sig)))
+	buf.Write(sig)
+	buf.Write(setBytes)
+
+	enc := base32.StdEncoding.WithPadding(base32.NoPadding)
+	return enc.EncodeToString(buf.Bytes())
+}
+
 type DownloadSetItem struct {
 	Kind   string
 	Name   string
@@ -138,32 +160,8 @@ func wsSignDownloadSet(req *restful.Request, resp *restful.Response) {
 		Items:  setReq.Items,
 	}
 
-	buf := new(bytes.Buffer)
-	{
-		setBytes := []byte(set.Encode())
-
-		w := lz4.NewWriter(buf)
-		w.Write(setBytes)
-		w.Close()
-	}
-
-	setBytes := buf.Bytes()
-
-	privkey, pubkey := dlsSigningKeys()
-	sig := ed25519.Sign(privkey, setBytes)
-
-	if !ed25519.Verify(pubkey, setBytes, sig) {
-		wsError(resp, fmt.Errorf("signature self-check failed"))
-		return
-	}
-
-	buf = bytes.NewBuffer(make([]byte, 0, 1+len(sig)+len(setBytes)))
-	buf.WriteByte(byte(len(sig)))
-	buf.Write(sig)
-	buf.Write(setBytes)
-
-	enc := base32.StdEncoding.WithPadding(base32.NoPadding)
-	resp.WriteEntity(enc.EncodeToString(buf.Bytes()))
+	privKey, _ := dlsSigningKeys()
+	resp.WriteEntity(set.Signed(privKey))
 }
 
 func getDlSet(req *restful.Request) (*DownloadSet, *httperr.Error) {