guru
/
squid-rewriter
1
0
Fork 0
Code Issues Releases Activity

Add flag to log to file
ci/woodpecker/push/lint Pipeline was successful Details
ci/woodpecker/push/test Pipeline was successful Details

This commit is contained in:
Maurizio Porrato 2022-09-21 07:34:40 +01:00
parent c955c8248b
commit e96566a724
1 changed files with 27 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
main.go
View File

@ -17,9 +17,12 @@ import (
var Flags struct {
ConfigFile string
LogFile string
Verbose bool
}
var L *log.Logger
func RewritesFromConfig(cfgfile string) (*trie.Trie, error) {
t := trie.NewTrie()
cfg, err := config.Load(cfgfile)
@ -75,7 +78,7 @@ func RewritesFromConfig(cfgfile string) (*trie.Trie, error) {
if d == nil {
continue
}
log.Printf("Fetching %s mirrors %v", distroName, repoList)
L.Printf("Fetching %s mirrors %v", distroName, repoList)
wg.Add(1)
go d.FetchMirrors(repoList, repoMirrors, &wg)
}
@ -83,7 +86,7 @@ func RewritesFromConfig(cfgfile string) (*trie.Trie, error) {
wg.Wait()
close(repoMirrors)
if len(repoMirrors) > 0 {
log.Println("Waiting for all repo rewrite rules to be processed...")
L.Println("Waiting for all repo rewrite rules to be processed...")
for len(repoMirrors) > 0 {
time.Sleep(5 * time.Millisecond)
}
@ -96,7 +99,7 @@ func RewritesFromConfig(cfgfile string) (*trie.Trie, error) {
}
if len(r.Urls) > 0 {
if Flags.Verbose {
log.Printf("Loading static rewites for %s", r.Name)
L.Printf("Loading static rewites for %s", r.Name)
}
for _, url := range r.Urls {
t.Put(url, r.Destination)
@ -104,7 +107,7 @@ func RewritesFromConfig(cfgfile string) (*trie.Trie, error) {
}
if r.Filename != "" {
if Flags.Verbose {
log.Printf("Loading rewrites from %s for %s", r.Filename, r.Name)
L.Printf("Loading rewrites from %s for %s", r.Filename, r.Name)
}
f, err := os.Open(r.Filename)
if err == nil {
@ -119,7 +122,7 @@ func RewritesFromConfig(cfgfile string) (*trie.Trie, error) {
}
}
} else {
log.Printf("Can't open file %s", r.Filename)
L.Printf("Can't open file %s", r.Filename)
}
}
}
@ -128,7 +131,7 @@ func RewritesFromConfig(cfgfile string) (*trie.Trie, error) {
}
func Rewrite(t *trie.Trie) {
log.Println("Listening for requests")
L.Println("Listening for requests")
scanner := bufio.NewScanner(os.Stdin)
for scanner.Scan() {
req := squid.ParseRequest(scanner.Text())
@ -143,7 +146,7 @@ func Rewrite(t *trie.Trie) {
newUrl := sdest + req.Url[len(prefix):]
resp.RewriteTo(newUrl)
if Flags.Verbose {
log.Printf("Rewriting %s to %s", req.Url, newUrl)
L.Printf("Rewriting %s to %s", req.Url, newUrl)
}
}
} else {
@ -151,11 +154,12 @@ func Rewrite(t *trie.Trie) {
}
os.Stdout.WriteString(resp.Format() + "\n")
}
log.Println("End of input stream. Exiting.")
L.Println("End of input stream. Exiting.")
}
func init() {
flag.StringVar(&Flags.ConfigFile, "c", "", "Path to rewrite config file")
flag.StringVar(&Flags.LogFile, "l", "", "Path to log file")
flag.BoolVar(&Flags.Verbose, "v", false, "Verbose logging")
}
@ -169,12 +173,12 @@ func LoadRewrites() *trie.Trie {
if Flags.ConfigFile != "" {
t, err = RewritesFromConfig(Flags.ConfigFile)
if err != nil {
log.Fatalf("Can't load the specified rewrite file: %s", Flags.ConfigFile)
L.Fatalf("Can't load the specified rewrite file: %s", Flags.ConfigFile)
}
} else {
for _, filename := range RewriteFileNames {
if Flags.Verbose {
log.Printf("Trying to load rewrites from %s", filename)
L.Printf("Trying to load rewrites from %s", filename)
}
t, err = RewritesFromConfig(filename)
if err == nil {
@ -182,16 +186,27 @@ func LoadRewrites() *trie.Trie {
}
}
if err != nil {
log.Fatalf("Can't load any of the predefined rewrite files: %v", RewriteFileNames)
L.Fatalf("Can't load any of the predefined rewrite files: %v", RewriteFileNames)
}
}
elapsed := time.Since(start)
log.Printf("Loaded %d unique rewrites in %s", t.Count(), elapsed)
L.Printf("Loaded %d unique rewrites in %s", t.Count(), elapsed)
return t
}
func main() {
flag.Parse()
if Flags.LogFile != "" {
file, err := os.OpenFile(Flags.LogFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
L.Fatal(err)
}
defer file.Close()
L = log.New(file, "", log.LstdFlags)
} else {
L = log.Default()
}
Rewrite(LoadRewrites())
}