add hosts file support

This commit is contained in:
Leo Antunes 2019-03-26 23:26:29 +01:00
parent b1ff41c523
commit e89a7141e5
5 changed files with 308 additions and 0 deletions

View File

@ -184,6 +184,7 @@ func (c *cluster) members() (<-chan []node, <-chan error) {
continue continue
} }
nodes = append(nodes, node{ nodes = append(nodes, node{
Name: n.Name,
Addr: n.Addr, Addr: n.Addr,
nodeMeta: meta, nodeMeta: meta,
}) })
@ -200,6 +201,7 @@ func (c *cluster) members() (<-chan []node, <-chan error) {
} }
type node struct { type node struct {
Name string
Addr net.IP Addr net.IP
nodeMeta nodeMeta
} }

View File

@ -17,6 +17,7 @@ type config struct {
WireguardPort int `id:"wireguard-port" desc:"port used for wireguard traffic (UDP); must be the same across cluster" default:"51820"` WireguardPort int `id:"wireguard-port" desc:"port used for wireguard traffic (UDP); must be the same across cluster" default:"51820"`
OverlayNet *network `id:"overlay-net" desc:"the network in which to allocate addresses for the overlay mesh network (CIDR format); smaller networks increase the chance of IP collision" default:"10.0.0.0/8"` OverlayNet *network `id:"overlay-net" desc:"the network in which to allocate addresses for the overlay mesh network (CIDR format); smaller networks increase the chance of IP collision" default:"10.0.0.0/8"`
Interface string `desc:"name of the wireguard interface to create and manage" default:"wgoverlay"` Interface string `desc:"name of the wireguard interface to create and manage" default:"wgoverlay"`
NoEtcHosts bool `id:"no-etc-hosts" desc:"disable writing of entries to /etc/hosts"`
LogLevel string `id:"log-level" desc:"set the verbosity (debug/info/warn/error)" default:"warn"` LogLevel string `id:"log-level" desc:"set the verbosity (debug/info/warn/error)" default:"warn"`
// for easier local testing // for easier local testing

164
etchosts/etchosts.go Normal file
View File

@ -0,0 +1,164 @@
package etchosts
import (
"bufio"
"fmt"
"io"
"io/ioutil"
"os"
"path"
"strings"
log "github.com/sirupsen/logrus"
)
// DefaultBanner is the default magic comment used to identify entries managed by etchosts
const DefaultBanner = "# ! MANAGED AUTOMATICALLY !"
// DefaultPath is the default path used to write hosts entries
const DefaultPath = "/etc/hosts"
// EtcHosts contains the options used to write hosts entries.
// The zero value can be used to write to DefaultPath using DefaultBanner as a marker.
type EtcHosts struct {
// Banner is the magic comment used to identify entries managed by etchosts; if not set, will use DefaultBanner.
// It must start with "#" to mark it as a comment.
Banner string
// Path is the path to the /etc/hosts file; if not set, will use DefaultPath.
Path string
// Logger is an optional logrus.StdLogger interface, used for debugging.
Logger log.StdLogger
}
// WriteEntries is used to write the hosts entries to EtcHosts.Path
// Each IP address with their (potentially multiple) hostnames are written to a line marked with EtcHosts.Banner, to
// avoid overwriting preexisting entries.
func (eh *EtcHosts) WriteEntries(ipsToNames map[string][]string) error {
hostsPath := eh.Path
if hostsPath == "" {
hostsPath = DefaultPath
}
// We do not want to create the hosts file; if it's not there, we probably have the wrong path.
etcHosts, err := os.OpenFile(hostsPath, os.O_RDWR, 0644)
if err != nil {
return fmt.Errorf("could not open %s for reading: %s", hostsPath, err)
}
defer etcHosts.Close()
// create tmpfile in same folder as
tmp, err := ioutil.TempFile(path.Dir(hostsPath), "etchosts")
if err != nil {
return fmt.Errorf("could not create tempfile")
}
// remove tempfile; this might fail if we managed to move it, which is ok
defer func(file *os.File) {
file.Close()
if err := os.Remove(file.Name()); err != nil && !os.IsNotExist(err) {
if eh.Logger != nil {
eh.Logger.Printf("unexpected error trying to remove temp file %s: %s", file.Name(), err)
}
}
}(tmp)
if err := eh.writeEntries(etcHosts, tmp, ipsToNames); err != nil {
return err
}
if err = eh.movePreservePerms(tmp, etcHosts); err != nil {
return err
}
return nil
}
func (eh *EtcHosts) writeEntries(orig io.Reader, dest io.Writer, ipsToNames map[string][]string) error {
banner := eh.Banner
if banner == "" {
banner = DefaultBanner
}
// go through file and update existing entries/prune nonexistent entries
scanner := bufio.NewScanner(orig)
for scanner.Scan() {
line := scanner.Text()
if strings.HasSuffix(strings.TrimSpace(line), strings.TrimSpace(banner)) {
tokens := strings.Fields(line)
if len(tokens) < 1 {
continue // remove empty managed line
}
ip := tokens[0]
if names, ok := ipsToNames[ip]; ok {
err := eh.writeEntryWithBanner(dest, banner, ip, names)
if err != nil {
return err
}
delete(ipsToNames, ip) // otherwise we'll append it again below
}
} else {
// keep original unmanaged line
fmt.Fprintf(dest, "%s\n", line)
}
}
if err := scanner.Err(); err != nil {
return fmt.Errorf("error reading hosts file: %s", err)
}
// append remaining entries to file
for ip, names := range ipsToNames {
if err := eh.writeEntryWithBanner(dest, banner, ip, names); err != nil {
return err
}
}
return nil
}
func (eh *EtcHosts) writeEntryWithBanner(tmp io.Writer, banner, ip string, names []string) error {
if ip != "" && len(names) > 0 {
if eh.Logger != nil {
eh.Logger.Printf("writing entry for %s (%s)", ip, names)
}
if _, err := fmt.Fprintf(tmp, "%s\t%s\t%s\n", ip, strings.Join(names, " "), banner); err != nil {
return fmt.Errorf("error writing entry for %s: %s", ip, err)
}
}
return nil
}
func (eh *EtcHosts) movePreservePerms(src, dst *os.File) error {
if err := src.Sync(); err != nil {
return fmt.Errorf("could not sync changes to %s: %s", src.Name(), err)
}
etcHostsInfo, err := dst.Stat()
if err != nil {
return fmt.Errorf("could not stat %s: %s", dst.Name(), err)
}
if err = os.Rename(src.Name(), dst.Name()); err != nil {
log.Infof("could not rename to %s; falling back to copy (%s)", dst.Name(), err)
if _, err := src.Seek(0, io.SeekStart); err != nil {
return err
}
if _, err := dst.Seek(0, io.SeekStart); err != nil {
return err
}
if err := dst.Truncate(0); err != nil {
return err
}
_, err = io.Copy(dst, src)
return err
}
// ensure we're not running with some umask that might break things
if err := src.Chmod(etcHostsInfo.Mode()); err != nil {
return fmt.Errorf("could not chmod %s: %s", src.Name(), err)
}
// TODO: also keep user?
return nil
}

120
etchosts/etchosts_test.go Normal file
View File

@ -0,0 +1,120 @@
package etchosts
import (
"bytes"
"fmt"
"io"
"strings"
"testing"
log "github.com/sirupsen/logrus"
)
func TestEtcHosts_writeEntryWithBanner(t *testing.T) {
type args struct {
banner string
ip string
names []string
}
eh := &EtcHosts{}
tests := []struct {
name string
args args
wantTmp string
wantErr bool
}{
{"do not write empty ip", args{DefaultBanner, "", []string{"somename", "someothername"}}, "", false},
{"do not write empty names", args{DefaultBanner, "1.2.3.4", []string{}}, "", false},
{"complete entry", args{DefaultBanner, "1.2.3.4", []string{"somename", "someothername"}}, fmt.Sprintf("1.2.3.4\tsomename someothername\t%s\n", DefaultBanner), false},
{"custom banner", args{"# somebanner", "1.2.3.4", []string{"somename", "someothername"}}, fmt.Sprintf("1.2.3.4\tsomename someothername\t%s\n", "# somebanner"), false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tmp := &bytes.Buffer{}
if err := eh.writeEntryWithBanner(tmp, tt.args.banner, tt.args.ip, tt.args.names); (err != nil) != tt.wantErr {
t.Errorf("writeEntryWithBanner() error = %v, wantErr %v", err, tt.wantErr)
return
}
if gotTmp := tmp.String(); gotTmp != tt.wantTmp {
t.Errorf("writeEntryWithBanner() got:\n%#v, want\n%#v", gotTmp, tt.wantTmp)
}
})
}
}
func TestEtcHosts_writeEntries(t *testing.T) {
type fields struct {
Banner string
Path string
Logger log.StdLogger
}
type args struct {
orig io.Reader
ipsToNames map[string][]string
}
tests := []struct {
name string
fields fields
args args
wantDest string
wantErr bool
}{
{
"simple empty write",
fields{},
args{strings.NewReader(""), map[string][]string{"1.2.3.4": []string{"foo", "bar"}}},
"1.2.3.4\tfoo bar\t# ! MANAGED AUTOMATICALLY !\n",
false,
},
{
"do not touch comments",
fields{},
args{strings.NewReader("# some comment\n"), map[string][]string{"1.2.3.4": []string{"foo", "bar"}}},
"# some comment\n1.2.3.4\tfoo bar\t# ! MANAGED AUTOMATICALLY !\n",
false,
},
{
"do not touch existing entries",
fields{},
args{strings.NewReader("4.3.2.1 hostname1 hostname2\n"), map[string][]string{"1.2.3.4": []string{"foo", "bar"}}},
"4.3.2.1 hostname1 hostname2\n1.2.3.4\tfoo bar\t# ! MANAGED AUTOMATICALLY !\n",
false,
},
{
"remove managed entry not in map",
fields{},
args{strings.NewReader("4.3.2.1 fooz baarz # ! MANAGED AUTOMATICALLY !\n"), map[string][]string{"1.2.3.4": []string{"foo", "bar"}}},
"1.2.3.4\tfoo bar\t# ! MANAGED AUTOMATICALLY !\n",
false,
},
{
"custom banner",
fields{Banner: "# somebanner"},
args{strings.NewReader(""), map[string][]string{"1.2.3.4": []string{"foo", "bar"}}},
"1.2.3.4\tfoo bar\t# somebanner\n",
false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
eh := &EtcHosts{
Banner: tt.fields.Banner,
Path: tt.fields.Path,
Logger: tt.fields.Logger,
}
dest := &bytes.Buffer{}
if err := eh.writeEntries(tt.args.orig, dest, tt.args.ipsToNames); (err != nil) != tt.wantErr {
t.Errorf("EtcHosts.writeEntries() error = %v, wantErr %v", err, tt.wantErr)
return
}
if gotDest := dest.String(); gotDest != tt.wantDest {
t.Errorf("EtcHosts.writeEntries() = '%#v', want '%#v'", gotDest, tt.wantDest)
}
})
}
}
// 1.2.3.4 foo bar # ! MANAGED AUTOMATICALLY !
// 1.2.3.4 foo bar # ! MANAGED AUTOMATICALLY !

21
main.go
View File

@ -6,6 +6,8 @@ import (
"syscall" "syscall"
"github.com/sirupsen/logrus" "github.com/sirupsen/logrus"
"github.com/costela/wesher/etchosts"
) )
func main() { func main() {
@ -52,12 +54,31 @@ func main() {
if err := wg.upInterface(); err != nil { if err := wg.upInterface(); err != nil {
logrus.Errorf("could not up interface: %s", err) logrus.Errorf("could not up interface: %s", err)
} }
if !config.NoEtcHosts {
if err := writeToEtcHosts(nodes); err != nil {
logrus.Errorf("could not write hosts entries: %s", err)
}
}
case errs := <-errc: case errs := <-errc:
logrus.Errorf("could not receive node info: %s", errs) logrus.Errorf("could not receive node info: %s", errs)
case <-incomingSigs: case <-incomingSigs:
logrus.Info("terminating...") logrus.Info("terminating...")
cluster.leave() cluster.leave()
if err := writeToEtcHosts(nil); err != nil {
logrus.Errorf("could not remove stale hosts entries: %s", err)
}
os.Exit(0) os.Exit(0)
} }
} }
} }
func writeToEtcHosts(nodes []node) error {
hosts := make(map[string][]string, len(nodes))
for _, n := range nodes {
hosts[n.OverlayAddr.String()] = []string{n.Name}
}
hostsFile := &etchosts.EtcHosts{
Logger: logrus.StandardLogger(),
}
return hostsFile.WriteEntries(hosts)
}