github.com/grafana/pyroscope@v1.18.0/pkg/test/integration/debuginfod_server.go (about)

     1  package integration
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"net"
     7  	"net/http"
     8  	"os"
     9  	"path/filepath"
    10  	"strings"
    11  )
    12  
    13  type TestDebuginfodServer struct {
    14  	server     *http.Server
    15  	debugFiles map[string]string
    16  	listener   net.Listener
    17  }
    18  
    19  func NewTestDebuginfodServer() (*TestDebuginfodServer, error) {
    20  	listener, err := net.Listen("tcp", "127.0.0.1:0")
    21  	if err != nil {
    22  		return nil, fmt.Errorf("failed to listen: %w", err)
    23  	}
    24  
    25  	s := &TestDebuginfodServer{
    26  		debugFiles: make(map[string]string),
    27  		listener:   listener,
    28  	}
    29  
    30  	mux := http.NewServeMux()
    31  	mux.HandleFunc("/buildid/", s.handleBuildID)
    32  
    33  	s.server = &http.Server{
    34  		Handler: mux,
    35  	}
    36  
    37  	return s, nil
    38  }
    39  
    40  func (s *TestDebuginfodServer) AddDebugFile(buildID, filePath string) {
    41  	s.debugFiles[buildID] = filePath
    42  }
    43  
    44  func (s *TestDebuginfodServer) URL() string {
    45  	return fmt.Sprintf("http://%s", s.listener.Addr().String())
    46  }
    47  
    48  func (s *TestDebuginfodServer) Start() error {
    49  	go func() {
    50  		_ = s.server.Serve(s.listener)
    51  	}()
    52  	return nil
    53  }
    54  
    55  func (s *TestDebuginfodServer) Stop() error {
    56  	return s.server.Close()
    57  }
    58  
    59  func (s *TestDebuginfodServer) handleBuildID(w http.ResponseWriter, r *http.Request) {
    60  	parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/buildid/"), "/")
    61  	if len(parts) != 2 || parts[1] != "debuginfo" {
    62  		http.Error(w, "Invalid path", http.StatusBadRequest)
    63  		return
    64  	}
    65  
    66  	buildID := parts[0]
    67  	filePath, ok := s.debugFiles[buildID]
    68  	if !ok {
    69  		http.Error(w, "Build ID not found", http.StatusNotFound)
    70  		return
    71  	}
    72  
    73  	file, err := os.Open(filePath)
    74  	if err != nil {
    75  		http.Error(w, fmt.Sprintf("Failed to open debug file: %v", err), http.StatusInternalServerError)
    76  		return
    77  	}
    78  	defer file.Close()
    79  
    80  	fileInfo, err := file.Stat()
    81  	if err != nil {
    82  		http.Error(w, fmt.Sprintf("Failed to stat debug file: %v", err), http.StatusInternalServerError)
    83  		return
    84  	}
    85  
    86  	w.Header().Set("Content-Type", "application/octet-stream")
    87  	w.Header().Set("Content-Length", fmt.Sprintf("%d", fileInfo.Size()))
    88  	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", filepath.Base(filePath)))
    89  
    90  	_, err = io.Copy(w, file)
    91  	if err != nil {
    92  		return
    93  	}
    94  }