diff --git a/cmd/hostpersister/main.go b/cmd/hostpersister/main.go index 87517c4..722e7c4 100644 --- a/cmd/hostpersister/main.go +++ b/cmd/hostpersister/main.go @@ -62,6 +62,9 @@ func main() { sigs := make(chan os.Signal, 1) signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM) + //Finished notifier: + finNotf := make(chan bool, 1) + z := time.Now().Sub(y) log.Printf("[Main] Took '%s' to fully initialize modules\n", z.String()) @@ -69,9 +72,19 @@ func main() { syncDur := time.Duration(syncTime) * time.Millisecond go func(exec *bool) { + executePersistence(hostsFile, sourceFile, overwriteMode) + sleep(syncDur, finNotf) for *exec { + err := hostsFile.ReadHostsFile() + if err != nil { + log.Fatalln("Failed to load HOSTS_FILE") + } + err = sourceFile.ReadHostsFile() + if err != nil { + log.Fatalln("Failed to load SOURCE_FILE") + } executePersistence(hostsFile, sourceFile, overwriteMode) - time.Sleep(syncDur) + sleep(syncDur, finNotf) } }(&exec) @@ -80,6 +93,7 @@ func main() { fmt.Printf("\n") *exec = false + finNotf <- true a := time.Now() log.Printf("[Main] Signalling program exit...\n") @@ -107,3 +121,12 @@ func executePersistence(hostsFile *hosts.File, sourceFile *hosts.File, ovrw bool log.Println("[Main] Error Writing Hosts File.") } } + +func sleep(d time.Duration, n chan bool) { + select { + case <-n: + break + case <-time.After(d): + break + } +} diff --git a/hosts/File.go b/hosts/File.go index bea0098..141d7a3 100644 --- a/hosts/File.go +++ b/hosts/File.go @@ -9,52 +9,14 @@ import ( const readBufferSize = 8192 func NewHostsFile(filePath string) (*File, error) { - theFile, err := os.Open(filePath) - if err != nil { - return nil, err + theHostFile := &File{ + filePath: filePath, } - defer theFile.Close() - var theEntries []Entry - var lenIn int - lineEnding := "" - theCBuffer := "" - theBuffer := make([]byte, readBufferSize) - for err == nil { - lenIn, err = theFile.Read(theBuffer) - if lenIn > 0 { - theCBuffer += string(theBuffer[:lenIn]) - if lineEnding == "" { - if strings.Contains(theCBuffer, "\r\n") { - lineEnding = "\r\n" - } else if strings.Contains(theCBuffer, "\r") { - lineEnding = "\r" - } else if strings.Contains(theCBuffer, "\n") { - lineEnding = "\n" - } - } - if lineEnding == "\r\n" { - strings.ReplaceAll(theCBuffer, "\r\n", "\n") - } else if lineEnding == "\r" { - strings.ReplaceAll(theCBuffer, "\r", "\n") - } - splt := strings.Split(theCBuffer, "\n") - for i := 0; i < len(splt)-1; i++ { - theEntries = append(theEntries, NewHostsEntry(splt[i])) - } - theCBuffer = splt[len(splt)-1] - } + err := theHostFile.ReadHostsFile() + if err == nil { + return theHostFile, nil } - if err != io.EOF { - return nil, err - } - if theCBuffer != "" { - theEntries = append(theEntries, NewHostsEntry(theCBuffer)) - } - return &File{ - filePath: filePath, - Entries: theEntries, - lineEnding: lineEnding, - }, nil + return nil, err } type File struct { @@ -63,6 +25,51 @@ type File struct { lineEnding string } +func (f *File) ReadHostsFile() error { + f.Entries = nil + theFile, err := os.Open(f.filePath) + if err != nil { + return err + } + defer theFile.Close() + var lenIn int + f.lineEnding = "" + theCBuffer := "" + theBuffer := make([]byte, readBufferSize) + for err == nil { + lenIn, err = theFile.Read(theBuffer) + if lenIn > 0 { + theCBuffer += string(theBuffer[:lenIn]) + if f.lineEnding == "" { + if strings.Contains(theCBuffer, "\r\n") { + f.lineEnding = "\r\n" + } else if strings.Contains(theCBuffer, "\r") { + f.lineEnding = "\r" + } else if strings.Contains(theCBuffer, "\n") { + f.lineEnding = "\n" + } + } + if f.lineEnding == "\r\n" { + strings.ReplaceAll(theCBuffer, "\r\n", "\n") + } else if f.lineEnding == "\r" { + strings.ReplaceAll(theCBuffer, "\r", "\n") + } + splt := strings.Split(theCBuffer, "\n") + for i := 0; i < len(splt)-1; i++ { + f.Entries = append(f.Entries, NewHostsEntry(splt[i])) + } + theCBuffer = splt[len(splt)-1] + } + } + if err != io.EOF { + return err + } + if theCBuffer != "" { + f.Entries = append(f.Entries, NewHostsEntry(theCBuffer)) + } + return nil +} + func (f File) WriteHostsFile() error { theFile, err := os.OpenFile(f.filePath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644) if err != nil {