121 lines
2.3 KiB
Go
121 lines
2.3 KiB
Go
/* SPDX-License-Identifier: GPL-2.0
|
|
*
|
|
* Copyright (C) 2021 Jason A. Donenfeld. All Rights Reserved.
|
|
*/
|
|
|
|
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"crypto/rand"
|
|
"encoding/base64"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"path"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type ShortLink struct {
|
|
mu sync.RWMutex
|
|
diskCachePath string
|
|
links map[string]string
|
|
dirty bool
|
|
cacher *time.Timer
|
|
}
|
|
|
|
func NewShortLink(diskCachePath string) *ShortLink {
|
|
sl := &ShortLink{
|
|
diskCachePath: diskCachePath,
|
|
links: make(map[string]string, 4096),
|
|
}
|
|
f, err := os.Open(sl.diskCachePath)
|
|
if err == nil {
|
|
defer f.Close()
|
|
scanner := bufio.NewScanner(f)
|
|
for scanner.Scan() {
|
|
parts := strings.SplitN(scanner.Text(), "\t", 2)
|
|
if len(parts) != 2 || len(parts[0]) == 0 || len(parts[1]) == 0 {
|
|
continue
|
|
}
|
|
sl.links[parts[0]] = parts[1]
|
|
}
|
|
}
|
|
return sl
|
|
}
|
|
|
|
func (sl *ShortLink) Set(k, v string, overwrite bool) bool {
|
|
sl.mu.Lock()
|
|
defer sl.mu.Unlock()
|
|
if wantV, ok := sl.links[k]; ok && (!overwrite || wantV == v) {
|
|
return false
|
|
}
|
|
sl.links[k] = v
|
|
if !sl.dirty {
|
|
sl.dirty = true
|
|
time.AfterFunc(time.Second*10, sl.SaveToDisk)
|
|
}
|
|
return true
|
|
}
|
|
|
|
func (sl *ShortLink) New(v string) string {
|
|
for {
|
|
var bytes [3]byte
|
|
_, err := rand.Read(bytes[:])
|
|
if err != nil {
|
|
log.Panicf("RNG is broken: %v", err)
|
|
}
|
|
k := base64.RawURLEncoding.EncodeToString(bytes[:])
|
|
sl.mu.RLock()
|
|
_, exists := sl.links[k]
|
|
sl.mu.RUnlock()
|
|
if !exists && sl.Set(k, v, false) {
|
|
return k
|
|
}
|
|
}
|
|
}
|
|
|
|
func (sl *ShortLink) HandleRequest(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Add("Server", "WurGurBoo/1.0")
|
|
target := "https://www.wireguard.com/"
|
|
wantL, key := path.Split(r.URL.Path)
|
|
if wantL == "/l/" {
|
|
sl.mu.RLock()
|
|
v, ok := sl.links[key]
|
|
sl.mu.RUnlock()
|
|
if ok {
|
|
target = v
|
|
}
|
|
}
|
|
http.Redirect(w, r, target, 302)
|
|
}
|
|
|
|
func (sl *ShortLink) SaveToDisk() {
|
|
sl.mu.Lock()
|
|
dirty := sl.dirty
|
|
sl.dirty = false
|
|
sl.mu.Unlock()
|
|
if !dirty {
|
|
return
|
|
}
|
|
f, err := os.Create(sl.diskCachePath + ".tmp")
|
|
if err != nil {
|
|
log.Printf("Unable to backup short link db: %v", err)
|
|
return
|
|
}
|
|
defer f.Close()
|
|
sl.mu.RLock()
|
|
for k, v := range sl.links {
|
|
fmt.Fprintf(f, "%s\t%s\n", k, v)
|
|
}
|
|
sl.mu.RUnlock()
|
|
err = os.Rename(f.Name(), sl.diskCachePath)
|
|
if err != nil {
|
|
log.Printf("Unable to backup short link db: %v", err)
|
|
return
|
|
}
|
|
}
|