Persist the printer token on connect (survive unclean restarts)

connector.go

@@ -18,6 +18,10 @@ var (
 	errFileTooLarge = errors.New("File is too large.")
 )
 
+// OnPrinterUpdate, if set, is called after a successful connect so the caller
+// can persist the printer's (refreshed) token to disk. Set by main().
+var OnPrinterUpdate func(*Printer)
+
 type Payload struct {
 	File      io.Reader
 	Name      string
@@ -126,6 +130,12 @@ func (c *connector) Upload(printer *Printer, payload *Payload) error {
 			}
 			defer h.Disconnect()
 
+			// Persist the (possibly refreshed) token so it survives restarts,
+			// including an unclean exit that skips the shutdown defer.
+			if OnPrinterUpdate != nil {
+				OnPrinterUpdate(printer)
+			}
+
 			if payload.Size > FILE_SIZE_MAX {
 				return errFileTooLarge
 			}

main.go

@@ -83,6 +83,16 @@ func main() {
 
 	var printer *Printer
 	ls := NewLocalStorage(KnownHosts)
+	// Persist the token to known hosts as soon as it is (re)obtained on connect,
+	// so it survives restarts even on an unclean exit (the shutdown defer below
+	// only runs on a graceful exit).
+	OnPrinterUpdate = func(p *Printer) {
+		if p == nil {
+			return
+		}
+		ls.Add(p)
+		_ = ls.Save()
+	}
 	defer func() {
 		if printer != nil {
 			// update printer's token
@@ -140,8 +150,9 @@ func main() {
 				printer = printers[0]
 			}
 		} else {
-			// directly to printer using ip/hostname
-			printer = &Printer{IP: Host}
+			// directly to printer using ip/hostname; keep an ID so the token gets
+			// persisted (ls.Add skips ID-less printers)
+			printer = &Printer{IP: Host, ID: Host}
 		}
 	}