Golang ui.Controller类代码示例

func versionCheck(s *State, ctl *ui.Controller) {
	check := func() {
		resp, err := http.Get(versionEndpoint)
		if err != nil {
			log.Warn("Failed to get version info %s: %v", versionEndpoint, err)
			return
		}
		defer resp.Body.Close()

		var payload struct {
			Client struct {
				Version string
			}
		}

		err = json.NewDecoder(resp.Body).Decode(&payload)
		if err != nil {
			log.Warn("Failed to read version info: %v", err)
			return
		}

		if payload.Client.Version != version.MajorMinor() {
			s.newVersion = payload.Client.Version
			ctl.Update(s)
		}
	}

	// check immediately and then at a set interval
	check()
	for _ = range time.Tick(versionCheckInterval) {
		check()
	}
}

/**
 * Establishes and manages a tunnel proxy connection with the server
 */
func proxy(proxyAddr string, s *State, ctl *ui.Controller) {
	start := time.Now()
	remoteConn, err := conn.Dial(proxyAddr, "pxy", tlsConfig)
	if err != nil {
		// XXX: What is the proper response here?
		// display something to the user?
		// retry?
		// reset control connection?
		log.Error("Failed to establish proxy connection: %v", err)
		return
	}

	defer remoteConn.Close()
	err = msg.WriteMsg(remoteConn, &msg.RegProxyMsg{Url: s.publicUrl})
	if err != nil {
		// XXX: What is the proper response here?
		// display something to the user?
		// retry?
		// reset control connection?
		log.Error("Failed to write RegProxyMsg: %v", err)
		return
	}

	localConn, err := conn.Dial(s.opts.localaddr, "prv", nil)
	if err != nil {
		remoteConn.Warn("Failed to open private leg %s: %v", s.opts.localaddr, err)
		badGatewayBody := fmt.Sprintf(BadGateway, s.publicUrl, s.opts.localaddr, s.opts.localaddr)
		remoteConn.Write([]byte(fmt.Sprintf(`HTTP/1.0 502 Bad Gateway
Content-Type: text/html
Content-Length: %d

%s`, len(badGatewayBody), badGatewayBody)))
		return
	}
	defer localConn.Close()

	m := s.metrics
	m.proxySetupTimer.Update(time.Since(start))
	m.connMeter.Mark(1)
	ctl.Update(s)
	m.connTimer.Time(func() {
		localConn := s.protocol.WrapConn(localConn)
		bytesIn, bytesOut := conn.Join(localConn, remoteConn)
		m.bytesIn.Update(bytesIn)
		m.bytesOut.Update(bytesOut)
		m.bytesInCount.Inc(bytesIn)
		m.bytesOutCount.Inc(bytesOut)
	})
	ctl.Update(s)
}

func reconnectingControl(s *State, ctl *ui.Controller) {
	// how long we should wait before we reconnect
	maxWait := 30 * time.Second
	wait := 1 * time.Second

	for {
		control(s, ctl)

		if s.status == "online" {
			wait = 1 * time.Second
		}

		log.Info("Waiting %d seconds before reconnecting", int(wait.Seconds()))
		time.Sleep(wait)
		// exponentially increase wait time
		wait = 2 * wait
		wait = time.Duration(math.Min(float64(wait), float64(maxWait)))
		s.status = "reconnecting"
		ctl.Update(s)
	}
}

func progressWatcher(s *State, ctl *ui.Controller, progress chan int, complete chan int) {
	for {
		select {
		case pct, ok := <-progress:
			if !ok {
				close(complete)
				return
			} else if pct == 100 {
				s.update = ui.UpdateInstalling
				ctl.Update(s)
				close(complete)
				return
			} else {
				if pct%25 == 0 {
					log.Info("Downloading update %d%% complete", pct)
				}
				s.update = ui.UpdateStatus(pct)
				ctl.Update(s)
			}
		}
	}
}

/**
 * Establishes and manages a tunnel control connection with the server
 */
func control(s *State, ctl *ui.Controller) {
	defer func() {
		if r := recover(); r != nil {
			log.Error("control recovering from failure %v", r)
		}
	}()

	// establish control channel
	conn, err := conn.Dial(s.opts.server, "ctl", tlsConfig)
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	// register with the server
	err = msg.WriteMsg(conn, &msg.RegMsg{
		Protocol:  s.opts.protocol,
		OS:        runtime.GOOS,
		HttpAuth:  s.opts.httpAuth,
		Hostname:  s.opts.hostname,
		Subdomain: s.opts.subdomain,
		ClientId:  s.id,
		Version:   version.Proto,
		MmVersion: version.MajorMinor(),
		User:      s.opts.authtoken,
	})

	if err != nil {
		panic(err)
	}

	// wait for the server to ack our register
	var regAck msg.RegAckMsg
	if err = msg.ReadMsgInto(conn, &regAck); err != nil {
		panic(err)
	}

	if regAck.Error != "" {
		emsg := fmt.Sprintf("Server failed to allocate tunnel: %s", regAck.Error)
		ctl.Cmds <- ui.CmdQuit{Message: emsg}
		return
	}

	// update UI state
	s.publicUrl = regAck.Url
	conn.Info("Tunnel established at %v", s.GetPublicUrl())
	s.status = "online"
	s.serverVersion = regAck.MmVersion
	ctl.Update(s)

	SaveAuthToken(s.opts.authtoken)

	// start the heartbeat
	lastPong := time.Now().UnixNano()
	go heartbeat(&lastPong, conn)

	// main control loop
	for {
		var m msg.Message
		if m, err = msg.ReadMsg(conn); err != nil {
			panic(err)
		}

		switch m.(type) {
		case *msg.ReqProxyMsg:
			go proxy(regAck.ProxyAddr, s, ctl)

		case *msg.PongMsg:
			atomic.StoreInt64(&lastPong, time.Now().UnixNano())
		}
	}
}

func autoUpdate(s *State, ctl *ui.Controller, token string) {
	tryAgain := true

	params := make(url.Values)
	params.Add("version", version.MajorMinor())
	params.Add("os", runtime.GOOS)
	params.Add("arch", runtime.GOARCH)
	params.Add("user", token)

	updateUrl := updateEndpoint + "?" + params.Encode()
	checkUrl := checkEndpoint + "?" + params.Encode()

	update := func() {
		log.Info("Checking for update")
		available, err := update.NewDownload(checkUrl).Check()
		if err != nil {
			log.Error("Error while checking for update: %v", err)
			return
		}

		if !available {
			log.Info("No update available")
			return
		}

		// stop trying after a single download attempt
		// XXX: improve this so the we can:
		// 1. safely update multiple times
		// 2. only retry after a network connection failure
		tryAgain = false

		download := update.NewDownload(updateUrl)
		downloadComplete := make(chan int)

		go progressWatcher(s, ctl, download.Progress, downloadComplete)

		log.Info("Trying to update . . .")
		err, errRecover := download.GetAndUpdate()
		<-downloadComplete

		if err != nil {
			// log error to console
			log.Error("Error while updating ngrok: %v", err)

			if errRecover != nil {
				log.Error("Error while recovering from failed ngrok update, your binary may be missing: %v", errRecover.Error())
				params.Add("errorRecover", errRecover.Error())
			}

			// log error to ngrok.com's servers for debugging purposes
			params.Add("error", err.Error())
			resp, reportErr := http.PostForm("https://dl.ngrok.com/update/error", params)
			if err != nil {
				log.Error("Error while reporting update error: %v, %v", err, reportErr)
			}
			resp.Body.Close()

			// tell the user to update manually
			s.update = ui.UpdateAvailable
		} else {
			if !download.Available {
				// this is the way the server tells us to update manually
				log.Info("Server wants us to update manually")
				s.update = ui.UpdateAvailable
			} else {
				// tell the user the update is ready
				log.Info("Update ready!")
				s.update = ui.UpdateReady
			}
		}

		ctl.Update(s)
		return
	}

	// try to update immediately and then at a set interval
	update()
	for _ = range time.Tick(updateCheckInterval) {
		if !tryAgain {
			break
		}
		update()
	}
}

func autoUpdate(s *State, ctl *ui.Controller, token string) {
	update := func() (updateSuccessful bool) {
		params := make(url.Values)
		params.Add("version", version.MajorMinor())
		params.Add("os", runtime.GOOS)
		params.Add("arch", runtime.GOARCH)

		download := update.NewDownload()
		downloadComplete := make(chan int)
		go func() {
			for {
				select {
				case progress, ok := <-download.Progress:
					if !ok {
						close(downloadComplete)
						return
					} else if progress == 100 {
						s.update = ui.UpdateInstalling
						ctl.Update(s)
						close(downloadComplete)
						return
					} else {
						if progress%25 == 0 {
							log.Info("Downloading update %d%% complete", progress)
						}
						s.update = ui.UpdateStatus(progress)
						ctl.Update(s)
					}
				}
			}
		}()

		log.Info("Checking for update")
		err := download.UpdateFromUrl(updateEndpoint + "?" + params.Encode())
		<-downloadComplete
		if err != nil {
			log.Error("Error while updating ngrok: %v", err)
			if download.Available {
				s.update = ui.UpdateError
			} else {
				s.update = ui.UpdateNone
			}

			// record the error to ngrok.com's servers for debugging purposes
			params.Add("error", err.Error())
			params.Add("user", token)
			resp, err := http.PostForm("https://dl.ngrok.com/update/error", params)
			if err != nil {
				log.Error("Error while reporting update error")
			}
			resp.Body.Close()
		} else {
			if download.Available {
				log.Info("Marked update ready")
				s.update = ui.UpdateReady
				updateSuccessful = true
			} else {
				log.Info("No update available at this time")
				s.update = ui.UpdateNone
			}
		}

		ctl.Update(s)
		return
	}

	// try to update immediately and then at a set interval
	update()
	for _ = range time.Tick(updateCheckInterval) {
		if update() {
			// stop trying to update if the update function is successful
			// XXX: improve this by trying to download versions newer than
			// the last one we downloaded
			return
		}
	}
}