pprofweb/pkg/pprofweb/pprofweb.go

package pprofweb

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/DataDog/gostackparse"
	"github.com/go-chi/chi/v5"
	"github.com/rs/xid"
	"github.com/spf13/afero"
	"tuxpa.in/a/pprofweb/pkg/goroutineinstance"
	"tuxpa.in/a/pprofweb/pkg/profileinstance"
)

type Config struct {
	Expire        time.Duration
	MaxUploadSize int
}

type Server struct {
	profiles   profileinstance.InstanceStore
	goroutines goroutineinstance.InstanceStore

	Config
}

func NewServer(fs afero.Fs, c Config) *Server {
	srv := afero.NewCacheOnReadFs(fs, afero.NewMemMapFs(), 24*time.Hour)
	s := &Server{
		profiles:   profileinstance.NewFileBasedInstanceStore(srv, c.Expire),
		goroutines: goroutineinstance.NewFileBasedInstanceStore(srv, c.Expire),
		Config:     c,
	}

	if s.MaxUploadSize == 0 {
		s.MaxUploadSize = 32 << 20
	}

	return s
}

// handler returns a handler that servers the pprof web UI.
func (s *Server) HandleHTTP() func(chi.Router) {
	return func(r chi.Router) {
		r.HandleFunc("/", rootHandler)
		r.HandleFunc("/upload", s.HandleUpload)
		r.Route("/profile", func(r chi.Router) {
			r.Route("/{xid}", func(r chi.Router) {
				r.HandleFunc("/", s.ServeProfileInstance)
				r.HandleFunc("/*", s.ServeProfileInstance)
			})
		})
		r.Route("/static/goroutine", func(r chi.Router) {
			r.Get("/bundle.js", func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "application/javascript")
				w.Write(goroutineinstance.BundleJsEmbed)
			})
			r.Get("/bundle.js.map", func(w http.ResponseWriter, r *http.Request) {
				w.Write(goroutineinstance.BundleJsMapEmbed)
			})

		})
		r.Route("/goroutine", func(r chi.Router) {
			r.Route("/{xid}", func(r chi.Router) {
				r.HandleFunc("/", s.ServeGoroutineInstance)
				r.HandleFunc("/*", s.ServeGoroutineInstance)
			})
		})
	}
}
func (s *Server) ServeGoroutineInstance(w http.ResponseWriter, r *http.Request) {
	err := s.serveGoroutineInstance(w, r)
	if err != nil {
		http.Error(w, "instance not found: "+err.Error(), http.StatusNotFound)
		return
	}
}

func (s *Server) serveGoroutineInstance(w http.ResponseWriter, r *http.Request) error {
	log.Printf("serveInstance %s %s", r.Method, r.URL.String())
	sid := chi.URLParam(r, "xid")
	if sid == "" {
		return errors.New("no id sent")
	}
	id, err := xid.FromString(sid)
	if err != nil {
		return errors.New("invalid id")
	}
	inst, err := s.goroutines.GetInstance(id)
	if err != nil {
		return errors.New("instance not found: " + err.Error())
	}
	inst.ServeHTTP(w, r)
	return nil
}

func (s *Server) ServeProfileInstance(w http.ResponseWriter, r *http.Request) {
	err := s.serveProfileInstance(w, r)
	if err != nil {
		http.Error(w, "instance not found: "+err.Error(), http.StatusNotFound)
		return
	}
}

func (s *Server) serveProfileInstance(w http.ResponseWriter, r *http.Request) error {
	log.Printf("serveInstance %s %s", r.Method, r.URL.String())
	sid := chi.URLParam(r, "xid")
	if sid == "" {
		return errors.New("no id sent")
	}
	id, err := xid.FromString(sid)
	if err != nil {
		return errors.New("invalid id")
	}
	inst, err := s.profiles.GetInstance(id)
	if err != nil {
		return errors.New("instance not found: " + err.Error())
	}
	inst.ServeHTTP(w, r)
	return nil
}

func (s *Server) HandleUpload(w http.ResponseWriter, r *http.Request) {
	log.Printf("uploadHandler %s %s", r.Method, r.URL.String())
	if r.Method != http.MethodPost {
		http.Error(w, "wrong method", http.StatusMethodNotAllowed)
		return
	}
	err := s.handleUpload(w, r)
	if err != nil {
		log.Printf("upload error: %s", err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

func (s *Server) handleUpload(w http.ResponseWriter, r *http.Request) error {
	if err := r.ParseMultipartForm(int64(s.MaxUploadSize)); err != nil {
		return err
	}
	uploadedFile, _, err := r.FormFile("file")
	if err != nil {
		return err
	}
	defer uploadedFile.Close()
	bts, err := io.ReadAll(uploadedFile)
	if err != nil {
		return err
	}

	// sanity check that allows us to be a little lazy.
	if len(bts) < 128 {
		return errors.New("profile too small")
	}

	var redirString string
	// check if it is a goroutine dump
	if bytes.HasPrefix(bts, []byte("goroutine")) {
		// goroutine dump. we should immediately try to parse it
		routines, errs := gostackparse.Parse(bytes.NewReader(bts))
		if errs != nil {
			return fmt.Errorf("error parsing goroutine dump: %w", err)
		}
		if len(routines) == 0 {
			return errors.New("no goroutines found")
		}
		instance, err := s.goroutines.NewInstance(routines)
		if err != nil {
			return err
		}
		redirString = "goroutine/" + instance.Id().String()
	} else {
		instance, err := s.profiles.NewInstance(bts)
		if err != nil {
			return err
		}
		redirString = "profile/" + instance.Id().String()
	}
	for _, h := range r.Header.Values("Accept") {
		if strings.Contains(h, "text/html") {
			http.Redirect(w, r, redirString, http.StatusSeeOther)
			return nil
		}
	}
	w.Write([]byte(redirString))
	return nil
}