tofu: Add NewHostsFile function

This commit is contained in:
Adnan Maolood 2021-01-14 16:54:38 -05:00
parent da3e9ac0fe
commit 6f11910dff

View File

@ -9,6 +9,7 @@ import (
"errors" "errors"
"fmt" "fmt"
"io" "io"
"os"
"strconv" "strconv"
"strings" "strings"
"sync" "sync"
@ -118,19 +119,33 @@ func (k *KnownHosts) TOFU(hostname string, cert *x509.Certificate) error {
return nil 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 { type HostWriter struct {
bw *bufio.Writer bw *bufio.Writer
cl io.Closer
mu sync.Mutex mu sync.Mutex
} }
// NewHostsWriter returns a new host writer that writes to the provided writer. // NewHostWriter returns a new host writer that writes to
func NewHostsWriter(w io.Writer) *HostWriter { // the provided io.WriteCloser.
func NewHostWriter(w io.WriteCloser) *HostWriter {
return &HostWriter{ return &HostWriter{
bw: bufio.NewWriter(w), 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. // WriteHost writes the host to the underlying io.Writer.
func (h *HostWriter) WriteHost(host Host) error { func (h *HostWriter) WriteHost(host Host) error {
h.mu.Lock() h.mu.Lock()
@ -140,11 +155,18 @@ func (h *HostWriter) WriteHost(host Host) error {
h.bw.WriteByte('\n') h.bw.WriteByte('\n')
if err := h.bw.Flush(); err != nil { 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 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. // Host represents a host entry with a fingerprint using a certain algorithm.
type Host struct { type Host struct {
Hostname string // hostname Hostname string // hostname