github.com/coreos/goproxy@v0.0.0-20190513173959-f8dc2d7ba04e/examples/goproxy-websockets/main.go (about)

     1  package main
     2  
     3  import (
     4  	"crypto/tls"
     5  	"github.com/ecordell/goproxy"
     6  	"github.com/gorilla/websocket"
     7  	"log"
     8  	"net/http"
     9  	"net/url"
    10  	"os"
    11  	"os/signal"
    12  	"sync"
    13  	"time"
    14  )
    15  
    16  var upgrader = websocket.Upgrader{} // use default options
    17  
    18  func echo(w http.ResponseWriter, r *http.Request) {
    19  	c, err := upgrader.Upgrade(w, r, nil)
    20  	if err != nil {
    21  		log.Print("upgrade:", err)
    22  		return
    23  	}
    24  	defer c.Close()
    25  	for {
    26  		mt, message, err := c.ReadMessage()
    27  		if err != nil {
    28  			log.Println("read:", err)
    29  			break
    30  		}
    31  		log.Printf("recv: %s", message)
    32  		err = c.WriteMessage(mt, message)
    33  		if err != nil {
    34  			log.Println("write:", err)
    35  			break
    36  		}
    37  	}
    38  }
    39  
    40  func StartEchoServer(wg *sync.WaitGroup) {
    41  	log.Println("Starting echo server")
    42  	wg.Add(1)
    43  	go func() {
    44  		http.HandleFunc("/", echo)
    45  		err := http.ListenAndServeTLS(":12345", "localhost.pem", "localhost-key.pem", nil)
    46  		if err != nil {
    47  			panic("ListenAndServe: " + err.Error())
    48  		}
    49  		wg.Done()
    50  	}()
    51  }
    52  
    53  func StartProxy(wg *sync.WaitGroup) {
    54  	log.Println("Starting proxy server")
    55  	wg.Add(1)
    56  	go func() {
    57  		proxy := goproxy.NewProxyHttpServer()
    58  		proxy.OnRequest().HandleConnect(goproxy.AlwaysMitm)
    59  		proxy.Verbose = true
    60  
    61  		err := http.ListenAndServe(":54321", proxy)
    62  		if err != nil {
    63  			log.Fatal(err.Error())
    64  		}
    65  		wg.Done()
    66  	}()
    67  }
    68  
    69  func main() {
    70  	interrupt := make(chan os.Signal, 1)
    71  	signal.Notify(interrupt, os.Interrupt)
    72  	wg := &sync.WaitGroup{}
    73  	StartEchoServer(wg)
    74  	StartProxy(wg)
    75  
    76  	endpointUrl := "wss://localhost:12345"
    77  	proxyUrl := "wss://localhost:54321"
    78  
    79  	surl, _ := url.Parse(proxyUrl)
    80  	dialer := websocket.Dialer{
    81  		Subprotocols:    []string{"p1"},
    82  		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
    83  		Proxy:           http.ProxyURL(surl),
    84  	}
    85  
    86  	c, _, err := dialer.Dial(endpointUrl, nil)
    87  	if err != nil {
    88  		log.Fatal("dial:", err)
    89  	}
    90  	defer c.Close()
    91  
    92  	done := make(chan struct{})
    93  
    94  	go func() {
    95  		defer c.Close()
    96  		defer close(done)
    97  		for {
    98  			_, message, err := c.ReadMessage()
    99  			if err != nil {
   100  				log.Println("read:", err)
   101  				return
   102  			}
   103  			log.Printf("recv: %s", message)
   104  		}
   105  	}()
   106  
   107  	ticker := time.NewTicker(time.Second)
   108  	defer ticker.Stop()
   109  
   110  	for {
   111  		select {
   112  		case t := <-ticker.C:
   113  			err := c.WriteMessage(websocket.TextMessage, []byte(t.String()))
   114  			if err != nil {
   115  				log.Println("write:", err)
   116  				return
   117  			}
   118  		case <-interrupt:
   119  			log.Println("interrupt")
   120  			// To cleanly close a connection, a client should send a close
   121  			// frame and wait for the server to close the connection.
   122  			err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
   123  			if err != nil {
   124  				log.Println("write close:", err)
   125  				return
   126  			}
   127  			select {
   128  			case <-done:
   129  			case <-time.After(time.Second):
   130  			}
   131  			c.Close()
   132  			return
   133  		}
   134  	}
   135  }