local-server/cmd/dkl-local-server/ws-downloads.go

package main

import (
	"crypto/rand"
	"encoding/base32"
	"log"
	"net/http"
	"strconv"
	"time"

	restful "github.com/emicklei/go-restful"
	"m.cluseau.fr/go/cow"
)

type DownloadSpec struct {
	Kind   string
	Name   string
	Assets []string

	createdAt time.Time
}

func wsAuthorizeDownload(req *restful.Request, resp *restful.Response) {
	var spec DownloadSpec

	if err := req.ReadEntity(&spec); err != nil {
		wsError(resp, err)
		return
	}

	if spec.Kind == "" || spec.Name == "" || len(spec.Assets) == 0 {
		resp.WriteErrorString(http.StatusBadRequest, "missing data")
		return
	}

	randBytes := make([]byte, 32)
	_, err := rand.Read(randBytes)
	if err != nil {
		wsError(resp, err)
		return
	}

	token := base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(randBytes)

	spec.createdAt = time.Now()

	wState.Change(func(v *State) {
		cow.MapSet(&v.Downloads, token, spec)
	})

	log.Printf("download token created for %s %q, assets %q", spec.Kind, spec.Name, spec.Assets)

	resp.WriteAsJson(token)
}

func wsDownload(req *restful.Request, resp *restful.Response) {
	token := req.PathParameter("token")
	asset := req.PathParameter("asset")

	if token == "" || asset == "" {
		wsNotFound(resp)
		return
	}

	var spec DownloadSpec
	found := false
	wState.Change(func(v *State) {
		var ok bool
		spec, ok = v.Downloads[token]
		if !ok {
			return
		}

		newAssets := make([]string, 0, len(spec.Assets))
		for _, a := range spec.Assets {
			if a == asset {
				found = true
			} else {
				newAssets = append(newAssets, a)
			}
		}

		if !found {
			wsNotFound(resp)
			return
		}

		cow.Map(&v.Downloads)

		if len(newAssets) == 0 {
			delete(v.Downloads, token)
		} else {
			spec.Assets = newAssets
			v.Downloads[token] = spec
		}
	})

	if !found {
		wsNotFound(resp)
		return
	}

	log.Printf("download via token: %s %q asset %q", spec.Kind, spec.Name, asset)

	cfg, err := readConfig()
	if err != nil {
		wsError(resp, err)
		return
	}

	setHeader := func(ext string) {
		resp.AddHeader("Content-Disposition", "attachment; filename="+strconv.Quote(spec.Kind+"_"+spec.Name+"_"+asset+ext))
	}

	switch spec.Kind {
	case "cluster":
		cluster := cfg.ClusterByName(spec.Name)
		if cluster == nil {
			wsNotFound(resp)
			return
		}

		switch asset {
		case "addons":
			setHeader(".yaml")
			resp.Write([]byte(cluster.Addons))

		default:
			wsNotFound(resp)
		}

	case "host":
		host := hostOrTemplate(cfg, spec.Name)
		if host == nil {
			wsNotFound(resp)
			return
		}

		switch asset {
		case "config", "bootstrap-config":
			setHeader(".yaml")
		default:
			setHeader("")
		}

		renderHost(resp.ResponseWriter, req.Request, asset, host, cfg)

	default:
		wsNotFound(resp)
	}
}