support for auth tokens, atexit fix

Author: petethepig

pkg/agent/profiler/profiler.go

@@ -12,6 +12,7 @@ import (
 type Config struct {
 	ApplicationName string // e.g backend.purchases
 	ServerAddress   string // e.g http://pyroscope.services.internal:4040
+	AuthToken       string
 }
 
 type Profiler struct {
@@ -21,6 +22,7 @@ type Profiler struct {
 // Start starts continuously profiling go code
 func Start(cfg Config) (*Profiler, error) {
 	u := remote.New(remote.RemoteConfig{
+		AuthToken:              cfg.AuthToken,
 		UpstreamAddress:        cfg.ServerAddress,
 		UpstreamThreads:        4,
 		UpstreamRequestTimeout: 30 * time.Second,

pkg/agent/upstream/remote/remote.go

@@ -31,6 +31,7 @@ type Remote struct {
 }
 
 type RemoteConfig struct {
+	AuthToken              string
 	UpstreamThreads        int
 	UpstreamAddress        string
 	UpstreamRequestTimeout time.Duration
@@ -99,7 +100,18 @@ func (u *Remote) uploadProfile(j *uploadJob) {
 	urlObj.RawQuery = q.Encode()
 	buf := j.t.Bytes()
 	log.Info("uploading at ", urlObj.String())
-	resp, err := u.client.Post(urlObj.String(), "binary/octet-stream+trie", bytes.NewReader(buf))
+
+	req, err := http.NewRequest("POST", urlObj.String(), bytes.NewReader(buf))
+	if err != nil {
+		log.Error("Error happened when uploading a profile:", err)
+		return
+	}
+	req.Header.Set("Content-Type", "binary/octet-stream+trie")
+	if u.cfg.AuthToken != "" {
+		req.Header.Set("Authorization", "Bearer "+u.cfg.AuthToken)
+	}
+	resp, err := u.client.Do(req)
+
 	if err != nil {
 		log.Error("Error happened when uploading a profile:", err)
 	}

pkg/config/config.go

@@ -22,6 +22,7 @@ type Agent struct {
 	AgentSpyName           string        `desc:"name of the spy you want to use"` // TODO: add options
 	AgentPID               int           `def:"-1" desc:"pid of the process you want to spy on"`
 	ServerAddress          string        `def:"http://localhost:4040" desc:"address of the pyroscope server"`
+	AuthToken              string        `def:"" desc:"authorization token used to upload profiling data"`
 	UpstreamThreads        int           `def:"4"`
 	UpstreamRequestTimeout time.Duration `def:"10s"`
 	UNIXSocketPath         string        `def:"<installPrefix>/var/run/pyroscope-agent.sock" desc:"path to a UNIX socket file"`
@@ -52,7 +53,7 @@ type Server struct {
 	MaxNodesSerialization int `def:"2048" desc:"max number of nodes used when saving profiles to disk"`
 	MaxNodesRender        int `def:"2048" desc:"max number of nodes used to display data on the frontend"`
 
-	// current only used in our demo app
+	// currently only used in our demo app
 	HideApplications []string `def:""`
 
 	AnalyticsOptOut bool `def:"false" desc:"disables analytics"`
@@ -79,6 +80,7 @@ type Exec struct {
 	DetectSubprocesses     bool          `def:"true" desc:"makes pyroscope keep track of and profile subprocesses of the main process"`
 	LogLevel               string        `def:"info", desc:"debug|info|warn|error"`
 	ServerAddress          string        `def:"http://localhost:4040" desc:"address of the pyroscope server"`
+	AuthToken              string        `def:"" desc:"authorization token used to upload profiling data"`
 	UpstreamThreads        int           `def:"4" desc:"number of upload threads"`
 	UpstreamRequestTimeout time.Duration `def:"10s" desc:"profile upload timeout"`
 	NoLogging              bool          `def:"false" desc:"disables logging from pyroscope"`

pkg/util/atexit/atexit.go

@@ -3,12 +3,15 @@ package atexit
 import (
 	"os"
 	"os/signal"
+	"sync"
 	"syscall"
 )
 
 var callbacks []func()
 
-func init() {
+var once sync.Once
+
+func initSignalHandler() {
 	signalChan := make(chan os.Signal, 1)
 	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
 	go func() {
@@ -20,5 +23,6 @@ func init() {
 }
 
 func Register(cb func()) {
+	once.Do(initSignalHandler)
 	callbacks = append(callbacks, cb)
 }