mirror of
https://github.com/openziti/zrok.git
synced 2024-12-26 08:38:50 +01:00
293 lines
8.4 KiB
Go
293 lines
8.4 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/base64"
|
|
"fmt"
|
|
"github.com/go-openapi/runtime"
|
|
httptransport "github.com/go-openapi/runtime/client"
|
|
"github.com/openziti/sdk-golang/ziti"
|
|
"github.com/openziti/sdk-golang/ziti/edge"
|
|
"github.com/openziti/zrok/environment"
|
|
"github.com/openziti/zrok/environment/env_core"
|
|
"github.com/openziti/zrok/rest_client_zrok"
|
|
"github.com/openziti/zrok/rest_client_zrok/share"
|
|
"github.com/openziti/zrok/rest_model_zrok"
|
|
"github.com/openziti/zrok/sdk/golang/sdk"
|
|
"github.com/openziti/zrok/tui"
|
|
"github.com/openziti/zrok/util"
|
|
"github.com/sirupsen/logrus"
|
|
"github.com/spf13/cobra"
|
|
"io"
|
|
"math/rand"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"syscall"
|
|
"time"
|
|
)
|
|
|
|
func init() {
|
|
loopCmd.AddCommand(newTestLoopPublicCommand().cmd)
|
|
}
|
|
|
|
type testLoopPublicCommand struct {
|
|
cmd *cobra.Command
|
|
loopers int
|
|
iterations int
|
|
statusEvery int
|
|
timeoutSeconds int
|
|
minPayload int
|
|
maxPayload int
|
|
minDwellMs int
|
|
maxDwellMs int
|
|
minPacingMs int
|
|
maxPacingMs int
|
|
frontendSelection []string
|
|
}
|
|
|
|
func newTestLoopPublicCommand() *testLoopPublicCommand {
|
|
cmd := &cobra.Command{
|
|
Use: "public",
|
|
Short: "Start a loop agent testing public proxy shares",
|
|
Args: cobra.ExactArgs(0),
|
|
}
|
|
r := &testLoopPublicCommand{cmd: cmd}
|
|
cmd.Run = r.run
|
|
cmd.Flags().IntVarP(&r.loopers, "loopers", "l", 1, "Number of current loopers to start")
|
|
cmd.Flags().IntVarP(&r.iterations, "iterations", "i", 1, "Number of iterations per looper")
|
|
cmd.Flags().IntVarP(&r.statusEvery, "status-every", "E", 100, "Show status every # iterations")
|
|
cmd.Flags().IntVarP(&r.timeoutSeconds, "timeout-seconds", "T", 30, "Time out after # seconds when sending http requests")
|
|
cmd.Flags().IntVar(&r.minPayload, "min-payload", 64, "Minimum payload size in bytes")
|
|
cmd.Flags().IntVar(&r.maxPayload, "max-payload", 10240, "Maximum payload size in bytes")
|
|
cmd.Flags().IntVar(&r.minDwellMs, "min-dwell-ms", 1000, "Minimum dwell time in milliseconds")
|
|
cmd.Flags().IntVar(&r.maxDwellMs, "max-dwell-ms", 1000, "Maximum dwell time in milliseconds")
|
|
cmd.Flags().IntVar(&r.minPacingMs, "min-pacing-ms", 0, "Minimum pacing in milliseconds")
|
|
cmd.Flags().IntVar(&r.maxPacingMs, "max-pacing-ms", 0, "Maximum pacing in milliseconds")
|
|
cmd.Flags().StringArrayVar(&r.frontendSelection, "frontends", []string{"public"}, "Selected frontends to use for the share")
|
|
return r
|
|
}
|
|
|
|
func (cmd *testLoopPublicCommand) run(_ *cobra.Command, _ []string) {
|
|
var loopers []*looper
|
|
for i := 0; i < cmd.loopers; i++ {
|
|
l := newLooper(i, cmd)
|
|
loopers = append(loopers, l)
|
|
go l.run()
|
|
}
|
|
c := make(chan os.Signal)
|
|
signal.Notify(c, os.Interrupt, syscall.SIGTERM)
|
|
go func() {
|
|
<-c
|
|
for _, looper := range loopers {
|
|
looper.stop = true
|
|
}
|
|
}()
|
|
for _, l := range loopers {
|
|
<-l.done
|
|
}
|
|
totalMismatches := 0
|
|
totalXfer := int64(0)
|
|
totalLoops := int64(0)
|
|
for _, l := range loopers {
|
|
deltaSeconds := l.stopTime.Sub(l.startTime).Seconds()
|
|
xfer := int64(float64(l.bytes) / deltaSeconds)
|
|
totalXfer += xfer
|
|
totalMismatches += l.mismatches
|
|
xferSec := util.BytesToSize(xfer)
|
|
totalLoops += l.loops
|
|
logrus.Infof("looper #%d: %d loops, %d mismatches, %s/sec", l.id, l.loops, l.mismatches, xferSec)
|
|
}
|
|
totalXferSec := util.BytesToSize(totalXfer)
|
|
logrus.Infof("total: %d loops, %d mismatches, %s/sec", totalLoops, totalMismatches, totalXferSec)
|
|
os.Exit(0)
|
|
}
|
|
|
|
type looper struct {
|
|
id int
|
|
cmd *testLoopPublicCommand
|
|
env *env_core.Environment
|
|
done chan struct{}
|
|
listener edge.Listener
|
|
zif string
|
|
zrok *rest_client_zrok.Zrok
|
|
shrToken string
|
|
proxyEndpoint string
|
|
auth runtime.ClientAuthInfoWriter
|
|
mismatches int
|
|
bytes int64
|
|
loops int64
|
|
startTime time.Time
|
|
stopTime time.Time
|
|
stop bool
|
|
}
|
|
|
|
func newLooper(id int, cmd *testLoopPublicCommand) *looper {
|
|
return &looper{
|
|
id: id,
|
|
cmd: cmd,
|
|
done: make(chan struct{}),
|
|
}
|
|
}
|
|
|
|
func (l *looper) run() {
|
|
defer close(l.done)
|
|
defer logrus.Infof("stopping #%d", l.id)
|
|
|
|
l.startup()
|
|
logrus.Infof("looper #%d, shrToken: %v, frontend: %v", l.id, l.shrToken, l.proxyEndpoint)
|
|
if l.serviceListener() {
|
|
l.dwell()
|
|
l.iterate()
|
|
}
|
|
logrus.Infof("looper #%d: complete", l.id)
|
|
l.shutdown()
|
|
}
|
|
|
|
func (l *looper) serviceListener() bool {
|
|
zcfg, err := ziti.NewConfigFromFile(l.zif)
|
|
if err != nil {
|
|
logrus.Errorf("error opening ziti config '%v': %v", l.zif, err)
|
|
return false
|
|
}
|
|
options := ziti.ListenOptions{
|
|
ConnectTimeout: 5 * time.Minute,
|
|
WaitForNEstablishedListeners: 1,
|
|
}
|
|
zctx, err := ziti.NewContext(zcfg)
|
|
if err != nil {
|
|
logrus.Errorf("error loading ziti context: %v", err)
|
|
return false
|
|
}
|
|
|
|
if l.listener, err = zctx.ListenWithOptions(l.shrToken, &options); err != nil {
|
|
logrus.Errorf("looper #%d, error listening: %v", l.id, err)
|
|
return false
|
|
}
|
|
|
|
go func() {
|
|
if err := http.Serve(l.listener, l); err != nil {
|
|
logrus.Errorf("looper #%d, error serving: %v", l.id, err)
|
|
}
|
|
}()
|
|
|
|
return true
|
|
}
|
|
|
|
func (l *looper) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
buf := new(bytes.Buffer)
|
|
io.Copy(buf, r.Body)
|
|
w.Write(buf.Bytes())
|
|
}
|
|
|
|
func (l *looper) startup() {
|
|
logrus.Infof("starting #%d", l.id)
|
|
|
|
env, err := environment.LoadRoot()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
if !env.IsEnabled() {
|
|
tui.Error("unable to load environment; did you 'zrok enable'?", nil)
|
|
}
|
|
l.env = env.Environment()
|
|
|
|
l.zif, err = env.ZitiIdentityNamed(env.EnvironmentIdentityName())
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
l.zrok, err = env.Client()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
l.auth = httptransport.APIKeyAuth("x-token", "header", l.env.Token)
|
|
tunnelReq := share.NewShareParams()
|
|
tunnelReq.Body = &rest_model_zrok.ShareRequest{
|
|
EnvZID: l.env.ZitiIdentity,
|
|
ShareMode: string(sdk.PublicShareMode),
|
|
FrontendSelection: l.cmd.frontendSelection,
|
|
BackendMode: string(sdk.ProxyBackendMode),
|
|
BackendProxyEndpoint: fmt.Sprintf("looper#%d", l.id),
|
|
AuthScheme: string(sdk.None),
|
|
}
|
|
tunnelReq.SetTimeout(60 * time.Second)
|
|
tunnelResp, err := l.zrok.Share.Share(tunnelReq, l.auth)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
l.shrToken = tunnelResp.Payload.ShrToken
|
|
l.proxyEndpoint = tunnelResp.Payload.FrontendProxyEndpoints[0]
|
|
}
|
|
|
|
func (l *looper) dwell() {
|
|
dwell := l.cmd.minDwellMs
|
|
if l.cmd.maxDwellMs-l.cmd.minDwellMs > 0 {
|
|
dwell = rand.Intn(l.cmd.maxDwellMs-l.cmd.minDwellMs) + l.cmd.minDwellMs
|
|
}
|
|
time.Sleep(time.Duration(dwell) * time.Millisecond)
|
|
}
|
|
|
|
func (l *looper) iterate() {
|
|
l.startTime = time.Now()
|
|
defer func() { l.stopTime = time.Now() }()
|
|
|
|
for i := 0; i < l.cmd.iterations && !l.stop; i++ {
|
|
if i > 0 && i%l.cmd.statusEvery == 0 {
|
|
logrus.Infof("looper #%d: iteration #%d", l.id, i)
|
|
}
|
|
sz := l.cmd.maxPayload
|
|
if l.cmd.maxPayload-l.cmd.minPayload > 0 {
|
|
sz = rand.Intn(l.cmd.maxPayload-l.cmd.minPayload) + l.cmd.minPayload
|
|
}
|
|
outpayload := make([]byte, sz)
|
|
outbase64 := base64.StdEncoding.EncodeToString(outpayload)
|
|
rand.Read(outpayload)
|
|
if req, err := http.NewRequest("POST", l.proxyEndpoint, bytes.NewBufferString(outbase64)); err == nil {
|
|
client := &http.Client{Timeout: time.Second * time.Duration(l.cmd.timeoutSeconds)}
|
|
if resp, err := client.Do(req); err == nil {
|
|
if resp.StatusCode != 200 {
|
|
logrus.Errorf("looper #%d unexpected response status code %v!", l.id, resp.StatusCode)
|
|
}
|
|
inpayload := new(bytes.Buffer)
|
|
io.Copy(inpayload, resp.Body)
|
|
inbase64 := inpayload.String()
|
|
if inbase64 != outbase64 {
|
|
logrus.Errorf("looper #%d payload mismatch!", l.id)
|
|
l.mismatches++
|
|
} else {
|
|
l.bytes += int64(len(outbase64))
|
|
logrus.Debugf("looper #%d payload match", l.id)
|
|
}
|
|
} else {
|
|
logrus.Errorf("looper #%d error: %v", l.id, err)
|
|
}
|
|
} else {
|
|
logrus.Errorf("looper #%d error creating request: %v", l.id, err)
|
|
}
|
|
pacingMs := l.cmd.maxPayload
|
|
if l.cmd.maxPacingMs-l.cmd.minPacingMs > 0 {
|
|
pacingMs = rand.Intn(l.cmd.maxPacingMs-l.cmd.minPacingMs) + l.cmd.minPacingMs
|
|
time.Sleep(time.Duration(pacingMs) * time.Millisecond)
|
|
}
|
|
l.loops++
|
|
}
|
|
}
|
|
|
|
func (l *looper) shutdown() {
|
|
if l.listener != nil {
|
|
if err := l.listener.Close(); err != nil {
|
|
logrus.Errorf("looper #%d error closing listener: %v", l.id, err)
|
|
}
|
|
}
|
|
|
|
untunnelReq := share.NewUnshareParams()
|
|
untunnelReq.Body = &rest_model_zrok.UnshareRequest{
|
|
EnvZID: l.env.ZitiIdentity,
|
|
ShrToken: l.shrToken,
|
|
}
|
|
if _, err := l.zrok.Share.Unshare(untunnelReq, l.auth); err != nil {
|
|
logrus.Errorf("error shutting down looper #%d: %v", l.id, err)
|
|
}
|
|
}
|