A HTTP Server to Download Docker Images with File Cache

docker-download-server.go

package main

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/sirupsen/logrus"
)

var (
	flagPort     = flag.Int("port", 8000, "HTTP server port.")
	flagHTTPRoot = flag.String("httpRoot", "/docker/", "HTTP serving root.")
	flagFileRoot = flag.String("fileRoot", "/tmp", "Root path of serving files.")
	compressor = "pzstd"
)

type DoingInfo struct {
	ch chan struct{}
	mu sync.Mutex

	Stage string
	Begin time.Time
}

type dockerHTTPHandler struct {
	mu    sync.Mutex
	root  string
	doing map[string]*DoingInfo
	chain http.Handler
}

func dockerFileServer(root string) http.Handler {
	return &dockerHTTPHandler{
		root:  root,
		doing: make(map[string]*DoingInfo),
		chain: http.FileServer(http.Dir(root)),
	}
}

func (h *dockerHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	//check if file exists, if not, run `docker pull && docker save`
	if _, err := os.Stat(path.Join(h.root, r.URL.Path)); !os.IsNotExist(err) {
		h.chain.ServeHTTP(w, r)
		return
	}

	if !strings.HasSuffix(r.URL.Path, ".tar.xz") {
		r.URL.Path += ".tar.xz"
		if _, err := os.Stat(path.Join(h.root, r.URL.Path)); !os.IsNotExist(err) {
			h.chain.ServeHTTP(w, r)
			return
		}
	}

	registry, uri := parseDockerURI(r.URL.Path)
	//validate path: [{registry}]/{repo}:{tag}
	logrus.WithField("image", uri).Info("Requested Docker image is not present. Try to pull it from registry.")
	if registry != "" {
		if err := os.MkdirAll(path.Join(h.root, registry), 0755); err != nil {
			logrus.WithError(err).Error("Mkdir failed.")
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
	h.mu.Lock()
	waiter := h.doing[uri]
	if waiter == nil {
		waiter = &DoingInfo{
			ch:    make(chan struct{}),
			Begin: time.Now(),
		}
		h.doing[uri] = waiter
		go h.dockerSave(uri, waiter)
	}
	h.mu.Unlock()

	if r.URL.Query().Get("block") == "" {
		w.WriteHeader(http.StatusServiceUnavailable)

		waiter.mu.Lock()
		text := fmt.Sprintf("Packing your docker images...\n * Currently status: %s\n * Time elapsed: %s\n\n\nPlease try it later...", waiter.Stage, time.Now().Sub(waiter.Begin).String())
		waiter.mu.Unlock()
		w.Write([]byte(text))
		return
	}

	<-waiter.ch
	h.chain.ServeHTTP(w, r)
}

func (h *dockerHTTPHandler) dockerSave(uri string, info *DoingInfo) {
	defer func() {
		close(info.ch)
		h.mu.Lock()
		delete(h.doing, uri)
		h.mu.Unlock()
	}()

	info.mu.Lock()
	info.Stage = "Pull Docker image from registry"
	info.mu.Unlock()

	logger := logrus.WithField("image", uri)
	logger.Info("Pull Docker image from registry.")
	cmdPull := exec.Command("docker", "pull", uri)
	cmdPull.Stdout, cmdPull.Stderr = os.Stdout, os.Stderr
	if err := cmdPull.Run(); err != nil {
		logger.WithError(err).Error("Docker pull failed.")
		return
	}
	info.mu.Lock()
	info.Stage = "Save Docker image on disk"
	info.mu.Unlock()

	dest := path.Join(h.root, uri+".tar.xz.tmp")
	logger.Info("Save Docker image on disk.")
	cmdSave := exec.Command("bash", "-c", fmt.Sprintf("docker save %s |%s > %s", uri, compressor, dest))
	cmdSave.Stdout, cmdSave.Stderr = os.Stdout, os.Stderr
	if err := cmdSave.Run(); err != nil {
		logger.WithError(err).Error("Docker save failed.")
		cmdRm := exec.Command("rm", dest)
		cmdRm.Stdout, cmdRm.Stderr = os.Stdout, os.Stderr
		if err := cmdRm.Run(); err != nil {
			logger.WithError(err).Error("Remove temp file failed.")
		}
		return
	}
	info.mu.Lock()
	info.Stage = "Clean up"
	info.mu.Unlock()

	cmdMv := exec.Command("mv", dest, strings.TrimSuffix(dest, ".tmp"))
	cmdMv.Stdout, cmdMv.Stderr = os.Stdout, os.Stderr
	if err := cmdMv.Run(); err != nil {
		logger.WithError(err).Error("Rename Docker image failed.")
		return
	}

	// logger.Info("Delete image from Docker daemon.")
	// cmdDel := exec.Command("docker", "rmi", uri)
	// cmdDel.Stdout, cmdDel.Stderr = os.Stdout, os.Stderr
	// if err := cmdDel.Run(); err != nil {
	//     logger.WithError(err).Error("Docker rmi failed.")
	//     return
	// }
}

func parseDockerURI(path string) (registry, uri string) {
	uri = strings.TrimLeft(path, "/")
	uri = strings.TrimSuffix(uri, ".tar.xz")
	if i := strings.LastIndex(uri, "/"); i >= 0 {
		registry = uri[:i]
	}
	return
}

func main() {
	flag.Parse()
	addr := fmt.Sprintf(":%d", *flagPort)

	logrus.WithField("port", *flagPort).Info("Start HTTP file server")
	http.Handle(*flagHTTPRoot, http.StripPrefix(*flagHTTPRoot, dockerFileServer(*flagFileRoot)))
	logrus.Fatal(http.ListenAndServe(addr, nil))
}