tofu: Add NewHostsFile function

1 files changed, 26 insertions, 4 deletions

diff --git a/tofu/tofu.go b/tofu/tofu.go
index e2f694a..9954734 100644
--- a/tofu/tofu.go
+++ b/tofu/tofu.go
@@ -9,6 +9,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
 	"strconv"
 	"strings"
 	"sync"
@@ -118,17 +119,31 @@ func (k *KnownHosts) TOFU(hostname string, cert *x509.Certificate) error {
 	return nil
 }
 
-// HostWriter writes host entries to an io.Writer.
+// HostWriter writes host entries to an io.WriteCloser.
+//
+// HostWriter is safe for concurrent use by multiple goroutines.
 type HostWriter struct {
 	bw *bufio.Writer
+	cl io.Closer
 	mu sync.Mutex
 }
 
-// NewHostsWriter returns a new host writer that writes to the provided writer.
-func NewHostsWriter(w io.Writer) *HostWriter {
+// NewHostWriter returns a new host writer that writes to
+// the provided io.WriteCloser.
+func NewHostWriter(w io.WriteCloser) *HostWriter {
 	return &HostWriter{
 		bw: bufio.NewWriter(w),
+		cl: w,
+	}
+}
+
+// NewHostsFile returns a new host writer that appends to the file at the given path.
+func NewHostsFile(path string) (*HostWriter, error) {
+	f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+	if err != nil {
+		return nil, err
 	}
+	return NewHostWriter(f), nil
 }
 
 // WriteHost writes the host to the underlying io.Writer.
@@ -140,11 +155,18 @@ func (h *HostWriter) WriteHost(host Host) error {
 	h.bw.WriteByte('\n')
 
 	if err := h.bw.Flush(); err != nil {
-		return fmt.Errorf("failed to write to hosts file: %w", err)
+		return fmt.Errorf("failed to write host: %w", err)
 	}
 	return nil
 }
 
+// Close closes the underlying io.WriteCloser.
+func (h *HostWriter) Close() error {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	return h.cl.Close()
+}
+
 // Host represents a host entry with a fingerprint using a certain algorithm.
 type Host struct {
 	Hostname    string      // hostname