diff --git a/README.md b/README.md index 29a8362..9ee2aee 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,7 @@ This server can be configured with these following parameters: | `HOME_REDIRECT` | (optional) which url to redirect when root url (`/`) is visited | `LISTEN_ADDR` | (optional) which network address to listen on (default `""` which means all interfaces) | | `PORT` | (optional) http port to listen on (default `8080`). +| `REDIRECT_STATUS` | (optional) HTTP status code to return upon redirect. (default `301`, Moved Permanently). ## Disclaimer diff --git a/main.go b/main.go index ce4cf57..558a524 100644 --- a/main.go +++ b/main.go @@ -7,6 +7,7 @@ import ( "net/http" "net/url" "os" + "strconv" "strings" "sync" "time" @@ -21,6 +22,7 @@ func main() { googleSheetsID := os.Getenv("GOOGLE_SHEET_ID") sheetName := os.Getenv("SHEET_NAME") homeRedirect := os.Getenv("HOME_REDIRECT") + redirectStatus := os.Getenv("REDIRECT_STATUS") ttlVal := os.Getenv("CACHE_TTL") ttl := time.Second * 5 @@ -32,16 +34,30 @@ func main() { ttl = v } - srv := &server{ - db: &cachedURLMap{ - ttl: ttl, - sheet: &sheetsProvider{ - googleSheetsID: googleSheetsID, - sheetName: sheetName, - }, + urlMap := &cachedURLMap{ + ttl: ttl, + sheet: &sheetsProvider{ + googleSheetsID: googleSheetsID, + sheetName: sheetName, }, + } + if err := urlMap.Init(); err != nil { + log.Fatalf("failed to initialize url map: %v", err) + } + + srv := &server{ + db: urlMap, homeRedirect: homeRedirect, } + if redirectStatus != "" { + s, err := strconv.Atoi(redirectStatus) + if err != nil { + log.Fatalf("failed to parse REDIRECT_STATUS as int: %v", err) + } + srv.redirectStatus = s + } else { + srv.redirectStatus = http.StatusMovedPermanently + } http.HandleFunc("/", srv.handler) @@ -52,11 +68,18 @@ func main() { } type server struct { - db *cachedURLMap - homeRedirect string + db *cachedURLMap + homeRedirect string + redirectStatus int } -type URLMap map[string]*url.URL +type mapData struct { + url *url.URL + hitCount int + rowIndex int +} + +type URLMap map[string]*mapData type cachedURLMap struct { sync.RWMutex @@ -67,7 +90,14 @@ type cachedURLMap struct { sheet *sheetsProvider } -func (c *cachedURLMap) Get(query string) (*url.URL, error) { +func (c *cachedURLMap) Init() error { + if err := c.sheet.Init(); err != nil { + return fmt.Errorf("failed to initialize sheet: %v", err) + } + return nil +} + +func (c *cachedURLMap) Get(query string) (*mapData, error) { if err := c.refresh(); err != nil { return nil, err } @@ -132,11 +162,12 @@ func (s *server) redirect(w http.ResponseWriter, req *http.Request) { } log.Printf("redirecting=%q to=%q", req.URL, redirTo.String()) - http.Redirect(w, req, redirTo.String(), http.StatusMovedPermanently) + http.Redirect(w, req, redirTo.String(), s.redirectStatus) + } func (s *server) findRedirect(req *url.URL) (*url.URL, error) { - path := strings.TrimPrefix(req.Path, "/") + path := strings.TrimPrefix(strings.ToLower(req.Path), "/") segments := strings.Split(path, "/") var discard []string @@ -147,7 +178,13 @@ func (s *server) findRedirect(req *url.URL) (*url.URL, error) { return nil, err } if v != nil { - return prepRedirect(v, strings.Join(discard, "/"), req.Query()), nil + go s.db.sheet.Write("C", v.rowIndex, + []interface{}{ + strconv.Itoa(v.hitCount + 1), + time.Now().Format(time.RFC3339), + }) + v.hitCount++ + return prepRedirect(v.url, strings.Join(discard, "/"), req.Query()), nil } discard = append([]string{segments[len(segments)-1]}, discard...) segments = segments[:len(segments)-1] @@ -173,7 +210,7 @@ func prepRedirect(base *url.URL, addPath string, query url.Values) *url.URL { func urlMap(in [][]interface{}) URLMap { out := make(URLMap) - for _, row := range in { + for i, row := range in { if len(row) < 2 { continue } @@ -185,6 +222,18 @@ func urlMap(in [][]interface{}) URLMap { if !ok || v == "" { continue } + hitCount := 0 + if len(row) >= 3 { + h, ok := row[2].(string) + if !ok || v == "" { + continue + } + hc, err := strconv.Atoi(h) + if err != nil { + log.Printf("warn: %s=%s hitCount invalid", k, h) + } + hitCount = hc + } k = strings.ToLower(k) u, err := url.Parse(v) @@ -197,7 +246,7 @@ func urlMap(in [][]interface{}) URLMap { if exists { log.Printf("warn: shortcut %q redeclared, overwriting", k) } - out[k] = u + out[k] = &mapData{u, hitCount, i + 1} } return out } diff --git a/sheetsprovider.go b/sheetsprovider.go index 6f65192..d95e95b 100644 --- a/sheetsprovider.go +++ b/sheetsprovider.go @@ -4,34 +4,63 @@ import ( "context" "fmt" "log" + "sync" "google.golang.org/api/sheets/v4" ) type sheetsProvider struct { + sync.RWMutex + client *sheets.Service googleSheetsID string sheetName string } -func (s *sheetsProvider) Query() ([][]interface{}, error) { +func (s *sheetsProvider) Init() error { if s.googleSheetsID == "" { - return nil, fmt.Errorf("GOOGLE_SHEET_ID not set") + return fmt.Errorf("GOOGLE_SHEET_ID not set") } srv, err := sheets.NewService(context.TODO()) if err != nil { - return nil, fmt.Errorf("unable to retrieve Sheets client: %v", err) + return fmt.Errorf("unable to retrieve Sheets client: %v", err) } + s.client = srv + return nil +} +func (s *sheetsProvider) Query() ([][]interface{}, error) { log.Println("querying sheet") - readRange := "A:B" + readRange := "A:D" if s.sheetName != "" { readRange = s.sheetName + "!" + readRange } - resp, err := srv.Spreadsheets.Values.Get(s.googleSheetsID, readRange).Do() + resp, err := s.client.Spreadsheets.Values.Get(s.googleSheetsID, readRange).Do() if err != nil { return nil, fmt.Errorf("unable to retrieve data from sheet: %v", err) } log.Printf("queried %d rows", len(resp.Values)) return resp.Values, nil } + +// Write will write the values rowwise, starting at the given column and row index. +func (s *sheetsProvider) Write(column string, rowIndex int, values []interface{}) error { + s.Lock() + defer s.Unlock() + log.Printf("writing %s to row %v", values, rowIndex) + writeRange := fmt.Sprintf("%s%d", column, rowIndex) + if s.sheetName != "" { + writeRange = s.sheetName + "!" + writeRange + } + _, err := s.client.Spreadsheets.Values.Update(s.googleSheetsID, writeRange, &sheets.ValueRange{ + Values: [][]interface{}{values}, + }).ValueInputOption("USER_ENTERED").Do() + if err != nil { + return fmt.Errorf("unable to write data to sheet: %v", err) + } + return nil +} + +func New() *sheetsProvider { + return &sheetsProvider{} +}