gloader/main.go

package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"

	"github.com/u-root/u-root/pkg/mount/gpt"
)

var (
	perm        = flag.String("perm", "/perm", "new val for perm")
	gloaderHome = "/perm/gloader"
)

func main() {
	flag.Parse()
	gloaderHome = filepath.Join(*perm, "gloader")
	os.MkdirAll(gloaderHome, 0777)
	err := mount()
	if err != nil {
		log.Println(err)
		os.Exit(1)
	}
	d := &Downloader{
		OnAdd:      save,
		OnComplete: save,
		OnDelete:   save,
		DataDir:    filepath.Join(gloaderHome, "data"),
	}
	loadQueue(d)
	d.Start("tcp", ":8844")
}

func save(d *Downloader, r Request) {
	var (
		content []byte
		err     error
	)
	content, err = json.Marshal(d.History.Queue)
	if err != nil {
		log.Println(err)
		return
	}
	err = ioutil.WriteFile(filepath.Join(gloaderHome, "history.json"), content, 0o666)
	if err != nil {
		log.Println(err)
		return
	}
	content, err = json.Marshal(d.Downloads.Queue)
	if err != nil {
		log.Println(err)
		return
	}
	err = ioutil.WriteFile(filepath.Join(gloaderHome, "queue.json"), content, 0o666)
	if err != nil {
		log.Println(err)
		return
	}
}

func loadQueue(d *Downloader) {
	var (
		f       io.ReadCloser
		err     error
		decoder *json.Decoder
	)
	f, err = os.Open(filepath.Join(gloaderHome, "history.json"))
	if err != nil {
		log.Println(err)
		return
	}
	decoder = json.NewDecoder(bufio.NewReader(f))
	err = decoder.Decode(&d.History.Queue)
	if err != nil {
		log.Println(err)
		return
	}
	f.Close()

	f, err = os.Open(filepath.Join(gloaderHome, "queue.json"))
	if err != nil {
		log.Println(err)
		return
	}
	decoder = json.NewDecoder(bufio.NewReader(f))
	err = decoder.Decode(&d.Downloads.Queue)
	if err != nil {
		log.Println(err)
		return
	}
	f.Close()
}

func mount() error {
	var (
		partUUIDb []byte
		partUUID  string
		dev       part
		err       error
		dataStat  os.FileInfo
		dataDir   = filepath.Join(gloaderHome, "data")
	)
	partUUIDb, err = ioutil.ReadFile(filepath.Join(gloaderHome, "partition"))
	partUUID = strings.TrimSpace(string(partUUIDb))
	if err != nil || partUUID == "" {
		return nil
	}
	dataStat, err = os.Stat(dataDir)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			err = os.MkdirAll(dataDir, 0o755)
			if err != nil {
				return fmt.Errorf("error creating datadir: %w", err)
			}
			dataStat, err = os.Stat(dataDir)
			if err != nil {
				return fmt.Errorf("error mounting datadir: %w", err)
			}
		} else {
			return fmt.Errorf("error mounting datadir: %w", err)
		}
	}
	if dataStat.IsDir() {
		var folder *os.File
		folder, err = os.Open(dataDir)
		if err != nil {
			return fmt.Errorf("error mounting datadir: %w", err)
		}
		_, err = folder.Readdir(1)
		if errors.Is(err, io.EOF) {
			log.Printf("mount %s %s\n", partUUID, dataDir)
			dev = findPartUUID(partUUID)
			err = syscall.Mount(dev.Path, dataDir, "ext4", 0, "")
			if err != nil {
				return fmt.Errorf("error mounting datadir: %w", err)
			}
			err = syscall.Mount(dev.Path, dataDir, "ext4", syscall.MS_SHARED|syscall.MS_REMOUNT, "")
			if err != nil {
				return fmt.Errorf("error remounting datadir: %w", err)
			}
			return nil
		}
		return fmt.Errorf("error mounting datadir: %w", err)
	}
	return fmt.Errorf("error mounting datadir: data dir %s is not a directory", dataDir)
}

type part struct {
	UUID string
	Path string
}

func findPartUUID(UUID string) part {
	var dev part
	err := filepath.Walk("/sys/block", func(path string, info os.FileInfo, err error) error {
		if err != nil {
			log.Printf("listPartUUID: %v", err)
			return nil
		}
		if info.Mode()&os.ModeSymlink == 0 {
			return nil
		}
		devname := "/dev/" + filepath.Base(path)
		f, err := os.Open(devname)
		if err != nil {
			log.Printf("listPartUUID: %v", err)
			return nil
		}
		defer f.Close()
		parttable, _ := gpt.New(f)
		if parttable.Primary != nil {
			for i, partition := range parttable.Primary.Parts {
				if partition.UniqueGUID.String() == UUID {
					dev = part{UUID: partition.UniqueGUID.String(), Path: devname + strconv.Itoa(i+1)}
				}
			}
		}

		return nil
	})
	if err != nil {
		return part{}
	}
	return dev
}