github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/scripts/traffic-duplicator/main.go (about)

     1  package main
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"context"
     7  	"flag"
     8  	"io"
     9  	"net/http"
    10  	"net/url"
    11  	"os/exec"
    12  	"regexp"
    13  	"strings"
    14  	"sync"
    15  	"time"
    16  
    17  	"github.com/sirupsen/logrus"
    18  )
    19  
    20  type target struct {
    21  	url     *url.URL
    22  	matcher *regexp.Regexp
    23  }
    24  
    25  var (
    26  	// input flags
    27  	bindAddr    string
    28  	logLevel    string
    29  	targetsPath string
    30  
    31  	// list of upstream targets
    32  	targetsMutex sync.RWMutex
    33  	targets      []target
    34  )
    35  
    36  func main() {
    37  	flag.StringVar(&bindAddr, "bind-addr", ":4040", "bind address for http server")
    38  	flag.StringVar(&logLevel, "log-level", "info", "log level")
    39  	flag.StringVar(&targetsPath, "targets-path", "./generate-targets.sh", "path to a script that generates upstream targets")
    40  	flag.Parse()
    41  
    42  	setupLogging()
    43  	go updateTargets()
    44  	startProxy()
    45  }
    46  
    47  func setupLogging() {
    48  	l, err := logrus.ParseLevel(logLevel)
    49  	if err != nil {
    50  		logrus.Fatal(err)
    51  	}
    52  	logrus.SetLevel(l)
    53  }
    54  
    55  func startProxy() {
    56  	logrus.WithFields(logrus.Fields{
    57  		"bind-addr":    bindAddr,
    58  		"targets-path": targetsPath,
    59  	}).Info("config")
    60  
    61  	err := http.ListenAndServe(bindAddr, http.HandlerFunc(handleConn))
    62  	logrus.WithError(err).WithField("bindAddr", bindAddr).Error("error listening")
    63  }
    64  
    65  func handleConn(_ http.ResponseWriter, r *http.Request) {
    66  	b, err := io.ReadAll(r.Body)
    67  	if err != nil {
    68  		logrus.WithError(err).Error("failed to read body")
    69  	}
    70  
    71  	targetsMutex.RLock()
    72  	targetsCopy := make([]target, len(targets))
    73  	copy(targetsCopy, targets)
    74  	targetsMutex.RUnlock()
    75  
    76  	for _, t := range targetsCopy {
    77  		appName := r.URL.Query().Get("name")
    78  		if t.matcher.MatchString(appName) {
    79  			logrus.WithField("target", t).Debug("uploading to upstream")
    80  			reader := bytes.NewReader(b)
    81  
    82  			r.URL.Scheme = t.url.Scheme
    83  			r.URL.Host = t.url.Host
    84  
    85  			resp, err := http.Post(r.URL.String(), r.Header.Get("Content-Type"), reader)
    86  			if err != nil {
    87  				logrus.WithError(err).WithField("target", t).Error("failed to upload to target")
    88  				continue
    89  			}
    90  			logrus.WithField("resp", resp).Debug("response")
    91  			if err := resp.Body.Close(); err != nil {
    92  				logrus.WithError(err).WithField("target", t).Error("failed to close response body")
    93  			}
    94  		}
    95  	}
    96  }
    97  
    98  func updateTargets() {
    99  	for {
   100  		ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
   101  		cmd := exec.CommandContext(ctx, targetsPath)
   102  		buf := bytes.Buffer{}
   103  		cmd.Stdout = &buf
   104  		err := cmd.Run()
   105  		if err != nil {
   106  			logrus.WithError(err).Error("failed to generate targets")
   107  		}
   108  		targetsMutex.Lock()
   109  		targets = generateTargets(bytes.NewReader(buf.Bytes()))
   110  		logrus.Debug("new targets:")
   111  		for _, t := range targets {
   112  			logrus.Debugf("* %s %s", t.url, t.matcher)
   113  		}
   114  		targetsMutex.Unlock()
   115  		time.Sleep(10 * time.Second)
   116  	}
   117  }
   118  
   119  func generateTargets(r io.Reader) []target {
   120  	var targets []target
   121  	scanner := bufio.NewScanner(r)
   122  	for scanner.Scan() {
   123  		line := strings.TrimSpace(scanner.Text())
   124  		arr := strings.SplitN(line, " ", 2)
   125  		if len(arr) < 2 {
   126  			arr = append(arr, ".*")
   127  		}
   128  
   129  		u, err := url.ParseRequestURI(arr[0])
   130  		if err != nil {
   131  			continue
   132  		}
   133  		matcher, err := regexp.Compile(arr[1])
   134  		if err != nil {
   135  			continue
   136  		}
   137  		targets = append(targets, target{
   138  			url:     u,
   139  			matcher: matcher,
   140  		})
   141  	}
   142  	return targets
   143  }