Skip to content

Instantly share code, notes, and snippets.

@cherrot
Last active June 27, 2022 17:55
Show Gist options
  • Save cherrot/3c4f45d9d0e4df5d57b98bcfe97f67dc to your computer and use it in GitHub Desktop.
Save cherrot/3c4f45d9d0e4df5d57b98bcfe97f67dc to your computer and use it in GitHub Desktop.
A HTTP Server to Download Docker Images with File Cache
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))
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment