github.com/MerlinKodo/quic-go@v0.39.2/example/main.go (about)

     1  package main
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"crypto/md5"
     7  	"errors"
     8  	"flag"
     9  	"fmt"
    10  	"io"
    11  	"log"
    12  	"mime/multipart"
    13  	"net/http"
    14  	"os"
    15  	"strconv"
    16  	"strings"
    17  	"sync"
    18  
    19  	_ "net/http/pprof"
    20  
    21  	"github.com/MerlinKodo/quic-go"
    22  	"github.com/MerlinKodo/quic-go/http3"
    23  	"github.com/MerlinKodo/quic-go/internal/testdata"
    24  	"github.com/MerlinKodo/quic-go/internal/utils"
    25  	"github.com/MerlinKodo/quic-go/logging"
    26  	"github.com/MerlinKodo/quic-go/qlog"
    27  )
    28  
    29  type binds []string
    30  
    31  func (b binds) String() string {
    32  	return strings.Join(b, ",")
    33  }
    34  
    35  func (b *binds) Set(v string) error {
    36  	*b = strings.Split(v, ",")
    37  	return nil
    38  }
    39  
    40  // Size is needed by the /demo/upload handler to determine the size of the uploaded file
    41  type Size interface {
    42  	Size() int64
    43  }
    44  
    45  // See https://en.wikipedia.org/wiki/Lehmer_random_number_generator
    46  func generatePRData(l int) []byte {
    47  	res := make([]byte, l)
    48  	seed := uint64(1)
    49  	for i := 0; i < l; i++ {
    50  		seed = seed * 48271 % 2147483647
    51  		res[i] = byte(seed)
    52  	}
    53  	return res
    54  }
    55  
    56  func setupHandler(www string) http.Handler {
    57  	mux := http.NewServeMux()
    58  
    59  	if len(www) > 0 {
    60  		mux.Handle("/", http.FileServer(http.Dir(www)))
    61  	} else {
    62  		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    63  			fmt.Printf("%#v\n", r)
    64  			const maxSize = 1 << 30 // 1 GB
    65  			num, err := strconv.ParseInt(strings.ReplaceAll(r.RequestURI, "/", ""), 10, 64)
    66  			if err != nil || num <= 0 || num > maxSize {
    67  				w.WriteHeader(400)
    68  				return
    69  			}
    70  			w.Write(generatePRData(int(num)))
    71  		})
    72  	}
    73  
    74  	mux.HandleFunc("/demo/tile", func(w http.ResponseWriter, r *http.Request) {
    75  		// Small 40x40 png
    76  		w.Write([]byte{
    77  			0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d,
    78  			0x49, 0x48, 0x44, 0x52, 0x00, 0x00, 0x00, 0x28, 0x00, 0x00, 0x00, 0x28,
    79  			0x01, 0x03, 0x00, 0x00, 0x00, 0xb6, 0x30, 0x2a, 0x2e, 0x00, 0x00, 0x00,
    80  			0x03, 0x50, 0x4c, 0x54, 0x45, 0x5a, 0xc3, 0x5a, 0xad, 0x38, 0xaa, 0xdb,
    81  			0x00, 0x00, 0x00, 0x0b, 0x49, 0x44, 0x41, 0x54, 0x78, 0x01, 0x63, 0x18,
    82  			0x61, 0x00, 0x00, 0x00, 0xf0, 0x00, 0x01, 0xe2, 0xb8, 0x75, 0x22, 0x00,
    83  			0x00, 0x00, 0x00, 0x49, 0x45, 0x4e, 0x44, 0xae, 0x42, 0x60, 0x82,
    84  		})
    85  	})
    86  
    87  	mux.HandleFunc("/demo/tiles", func(w http.ResponseWriter, r *http.Request) {
    88  		io.WriteString(w, "<html><head><style>img{width:40px;height:40px;}</style></head><body>")
    89  		for i := 0; i < 200; i++ {
    90  			fmt.Fprintf(w, `<img src="/demo/tile?cachebust=%d">`, i)
    91  		}
    92  		io.WriteString(w, "</body></html>")
    93  	})
    94  
    95  	mux.HandleFunc("/demo/echo", func(w http.ResponseWriter, r *http.Request) {
    96  		body, err := io.ReadAll(r.Body)
    97  		if err != nil {
    98  			fmt.Printf("error reading body while handling /echo: %s\n", err.Error())
    99  		}
   100  		w.Write(body)
   101  	})
   102  
   103  	// accept file uploads and return the MD5 of the uploaded file
   104  	// maximum accepted file size is 1 GB
   105  	mux.HandleFunc("/demo/upload", func(w http.ResponseWriter, r *http.Request) {
   106  		if r.Method == http.MethodPost {
   107  			err := r.ParseMultipartForm(1 << 30) // 1 GB
   108  			if err == nil {
   109  				var file multipart.File
   110  				file, _, err = r.FormFile("uploadfile")
   111  				if err == nil {
   112  					var size int64
   113  					if sizeInterface, ok := file.(Size); ok {
   114  						size = sizeInterface.Size()
   115  						b := make([]byte, size)
   116  						file.Read(b)
   117  						md5 := md5.Sum(b)
   118  						fmt.Fprintf(w, "%x", md5)
   119  						return
   120  					}
   121  					err = errors.New("couldn't get uploaded file size")
   122  				}
   123  			}
   124  			utils.DefaultLogger.Infof("Error receiving upload: %#v", err)
   125  		}
   126  		io.WriteString(w, `<html><body><form action="/demo/upload" method="post" enctype="multipart/form-data">
   127  				<input type="file" name="uploadfile"><br>
   128  				<input type="submit">
   129  			</form></body></html>`)
   130  	})
   131  
   132  	return mux
   133  }
   134  
   135  func main() {
   136  	// defer profile.Start().Stop()
   137  	go func() {
   138  		log.Println(http.ListenAndServe("localhost:6060", nil))
   139  	}()
   140  	// runtime.SetBlockProfileRate(1)
   141  
   142  	verbose := flag.Bool("v", false, "verbose")
   143  	bs := binds{}
   144  	flag.Var(&bs, "bind", "bind to")
   145  	www := flag.String("www", "", "www data")
   146  	tcp := flag.Bool("tcp", false, "also listen on TCP")
   147  	enableQlog := flag.Bool("qlog", false, "output a qlog (in the same directory)")
   148  	flag.Parse()
   149  
   150  	logger := utils.DefaultLogger
   151  
   152  	if *verbose {
   153  		logger.SetLogLevel(utils.LogLevelDebug)
   154  	} else {
   155  		logger.SetLogLevel(utils.LogLevelInfo)
   156  	}
   157  	logger.SetLogTimeFormat("")
   158  
   159  	if len(bs) == 0 {
   160  		bs = binds{"localhost:6121"}
   161  	}
   162  
   163  	handler := setupHandler(*www)
   164  	quicConf := &quic.Config{}
   165  	if *enableQlog {
   166  		quicConf.Tracer = func(ctx context.Context, p logging.Perspective, connID quic.ConnectionID) *logging.ConnectionTracer {
   167  			filename := fmt.Sprintf("server_%x.qlog", connID)
   168  			f, err := os.Create(filename)
   169  			if err != nil {
   170  				log.Fatal(err)
   171  			}
   172  			log.Printf("Creating qlog file %s.\n", filename)
   173  			return qlog.NewConnectionTracer(utils.NewBufferedWriteCloser(bufio.NewWriter(f), f), p, connID)
   174  		}
   175  	}
   176  
   177  	var wg sync.WaitGroup
   178  	wg.Add(len(bs))
   179  	for _, b := range bs {
   180  		bCap := b
   181  		go func() {
   182  			var err error
   183  			if *tcp {
   184  				certFile, keyFile := testdata.GetCertificatePaths()
   185  				err = http3.ListenAndServe(bCap, certFile, keyFile, handler)
   186  			} else {
   187  				server := http3.Server{
   188  					Handler:    handler,
   189  					Addr:       bCap,
   190  					QuicConfig: quicConf,
   191  				}
   192  				err = server.ListenAndServeTLS(testdata.GetCertificatePaths())
   193  			}
   194  			if err != nil {
   195  				fmt.Println(err)
   196  			}
   197  			wg.Done()
   198  		}()
   199  	}
   200  	wg.Wait()
   201  }