added cache

server/handler.go

@@ -20,29 +20,26 @@ import (
 	"fmt"
 	"net/http"
 	"os/exec"
+	"syscall"
+	"time"
+
+	"github.com/allegro/bigcache"
+	"github.com/sirupsen/logrus"
 )
 
 type Handler struct {
-	binary string
+	binary     string
+	baseParams []string
+	uid        int
+	logger     *logrus.Logger
+	cache      *bigcache.BigCache
 }
 
-func NewHandler(binary string) *Handler {
+func NewHandler(binary string, eviction time.Duration) *Handler {
 	h := new(Handler)
 	h.binary = binary
-	return h
-}
-
-func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	b, e := h.FetchDom(fmt.Sprintf("https://%s%s", r.Host, r.RequestURI))
-	if e != nil {
-		fmt.Fprintln(w, e.Error())
-		return
-	}
-	w.Write(b.Bytes())
-}
-
-func (h *Handler) FetchDom(url string) (*bytes.Buffer, error) {
-	params := []string{
+	h.uid = syscall.Getuid()
+	h.baseParams = []string{
 		"--disable-client-side-phishing-detection",
 		"--disable-cloud-import",
 		"--disable-default-apps",
@@ -63,15 +60,60 @@ func (h *Handler) FetchDom(url string) (*bytes.Buffer, error) {
 		"--headless",
 		"--disable-gpu",
 		"--disable-software-rasterizer",
-		"--no-sandbox",
 		"--dump-dom",
-		url,
 	}
-	cmd := exec.Command(h.binary, params...)
-	buf := bytes.NewBuffer(nil)
-	cmd.Stdout = buf
-	if e := cmd.Run(); e != nil {
-		return nil, e
+	h.logger = logrus.StandardLogger()
+	cacheConf := bigcache.DefaultConfig(eviction)
+	cacheConf.HardMaxCacheSize = 500 * 1024 * 1024 // 500MB
+	cache, e := bigcache.NewBigCache(cacheConf)
+	if e != nil {
+		h.logger.Fatal(e.Error())
+		return nil
 	}
-	return buf, nil
+	h.cache = cache
+	return h
+}
+
+func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	b, e := h.FetchDom(fmt.Sprintf("https://%s%s", r.Host, r.RequestURI))
+	if e != nil {
+		h.logger.Error(e)
+		if _, err := fmt.Fprintln(w, e.Error()); err != nil {
+			h.logger.Error(e)
+			return
+		}
+		return
+	}
+	if _, e := w.Write(b.Bytes()); e != nil {
+		h.logger.Error(e)
+	}
+
+}
+
+func (h *Handler) FetchDom(url string) (*bytes.Buffer, error) {
+	b, e := h.cache.Get(url)
+	if e != nil {
+		if e == bigcache.ErrEntryNotFound {
+			var params []string
+			if h.uid == 0 {
+				params = append(h.baseParams, "--no-sandbox")
+			} else {
+				params = h.baseParams
+			}
+			params = append(params, url)
+			cmd := exec.Command(h.binary, params...)
+			buf := bytes.NewBuffer(nil)
+			cmd.Stdout = buf
+			if e := cmd.Run(); e != nil {
+				return nil, e
+			}
+			if e := h.cache.Set(url, buf.Bytes()); e != nil {
+				h.logger.Error(e)
+			}
+			return buf, nil
+		} else {
+			return nil, e
+		}
+	}
+	return bytes.NewBuffer(b), nil
 }