Protocol Labs Code Sample

stratum.go

// Copyright (c) 2016 The Decred

// This is an excerpt from https://github.com/decred/gominer/blob/master/stratum/stratum.go
// It was introduced in:
// https://github.com/decred/gominer/pull/7
// https://github.com/decred/gominer/commit/55b4c6abb406f273b722d2c96621b6dfe95d3c5f

package stratum

import ()

// StratumConn starts the initial connection to a stratum pool and sets defaults
// in the pool object.
func StratumConn(pool, user, pass, proxy, proxyUser, proxyPass, version string) (*Stratum, error) {
	var stratum Stratum
	stratum.cfg.User = user
	stratum.cfg.Pass = pass
	stratum.cfg.Proxy = proxy
	stratum.cfg.ProxyUser = proxyUser
	stratum.cfg.ProxyPass = proxyPass
	stratum.cfg.Version = version

	log.Infof("Using pool: %v", pool)
	proto := "stratum+tcp://"
	if strings.HasPrefix(pool, proto) {
		pool = strings.Replace(pool, proto, "", 1)
	} else {
		err := errors.New("Only stratum pools supported.")
		return nil, err
	}
	var conn net.Conn
	var err error
	if stratum.cfg.Proxy != "" {
		proxy := &socks.Proxy{
			Addr:     stratum.cfg.Proxy,
			Username: stratum.cfg.ProxyUser,
			Password: stratum.cfg.ProxyPass,
		}
		conn, err = proxy.Dial("tcp", pool)
	} else {
		conn, err = net.Dial("tcp", pool)
	}
	if err != nil {
		return nil, err
	}
	stratum.ID = 1
	stratum.Conn = conn
	stratum.cfg.Pool = pool

	// We will set it for sure later but this really should be the value and
	// setting it here will prevent so incorrect matches based on the
	// default 0 value.
	stratum.authID = 2

	// Target for share is 1 unless we hear otherwise.
	stratum.Diff = 1
	stratum.Target, err = util.DiffToTarget(stratum.Diff, chainParams.PowLimit)
	if err != nil {
		return nil, err
	}
	stratum.PoolWork.NewWork = false
	stratum.Reader = bufio.NewReader(stratum.Conn)
	go stratum.Listen()

	err = stratum.Subscribe()
	if err != nil {
		return nil, err
	}
	err = stratum.Auth()
	if err != nil {
		return nil, err
	}

	return &stratum, nil
}

// Reconnect reconnects to a stratum server if the connection has been lost.
func (s *Stratum) Reconnect() error {
	var conn net.Conn
	var err error
	if s.cfg.Proxy != "" {
		proxy := &socks.Proxy{
			Addr:     s.cfg.Proxy,
			Username: s.cfg.ProxyUser,
			Password: s.cfg.ProxyPass,
		}
		conn, err = proxy.Dial("tcp", s.cfg.Pool)
	} else {
		conn, err = net.Dial("tcp", s.cfg.Pool)
	}
	if err != nil {
		return err
	}
	s.Conn = conn
	s.Reader = bufio.NewReader(s.Conn)
	err = s.Subscribe()
	if err != nil {
		return nil
	}
	// Should NOT need this.
	time.Sleep(5 * time.Second)
	// XXX Do I really need to re-auth here?
	err = s.Auth()
	if err != nil {
		return nil
	}
	return nil
}

// Listen is the listener for the incoming messages from the stratum pool.
func (s *Stratum) Listen() {
	log.Debug("Starting Listener")

	for {
		result, err := s.Reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				log.Error("Connection lost!  Reconnecting.")
				err = s.Reconnect()
				if err != nil {
					log.Error(err)
					log.Error("Reconnect failed.")
					os.Exit(1)
					return
				}

			} else {
				log.Error(err)
			}
			continue
		}

		log.Debug(strings.TrimSuffix(result, "\n"))
		resp, err := s.Unmarshal([]byte(result))
		if err != nil {
			log.Error(err)
			continue
		}

		switch resp.(type) {
		case *BasicReply:
			s.handleBasicReply(resp)
		case StratumMsg:
			s.handleStratumMsg(resp)
		case NotifyRes:
			s.handleNotifyRes(resp)
		case *SubscribeReply:
			s.handleSubscribeReply(resp)
		default:
			log.Info("Unhandled message: ", result)
		}
	}
}
func (s *Stratum) handleBasicReply(resp interface{}) {
	s.Lock()
	defer s.Unlock()
	aResp := resp.(*BasicReply)

	if int(aResp.ID.(uint64)) == int(s.authID) {
		if aResp.Result {
			log.Debug("Logged in")
		} else {
			log.Error("Auth failure.")
		}
	}
	if sliceContains(s.submitIDs, aResp.ID.(uint64)) {
		if aResp.Result {
			atomic.AddUint64(&s.ValidShares, 1)
			log.Debug("Share accepted")
		} else {
			atomic.AddUint64(&s.InvalidShares, 1)
			log.Error("Share rejected: ", aResp.Error.ErrStr)
		}
		s.submitIDs = sliceRemove(s.submitIDs, aResp.ID.(uint64))
	}
}

func (s *Stratum) handleStratumMsg(resp interface{}) {
	nResp := resp.(StratumMsg)
	log.Trace(nResp)
	// Too much is still handled in unmarshaler.  Need to
	// move stuff other than unmarshalling here.
	switch nResp.Method {
	case "client.show_message":
		log.Info(nResp.Params)
	case "client.reconnect":
		log.Debug("Reconnect requested")
		wait, err := strconv.Atoi(nResp.Params[2])
		if err != nil {
			log.Error(err)
			return
		}
		time.Sleep(time.Duration(wait) * time.Second)
		pool := nResp.Params[0] + ":" + nResp.Params[1]
		s.cfg.Pool = pool
		err = s.Reconnect()
		if err != nil {
			log.Error(err)
			// XXX should just die at this point
			// but we don't really have access to
			// the channel to end everything.
			return
		}

	case "client.get_version":
		log.Debug("get_version request received.")
		msg := StratumMsg{
			Method: nResp.Method,
			ID:     nResp.ID,
			Params: []string{"decred-gominer/" + s.cfg.Version},
		}
		m, err := json.Marshal(msg)
		if err != nil {
			log.Error(err)
			return
		}
		_, err = s.Conn.Write(m)
		if err != nil {
			log.Error(err)
			return
		}
		_, err = s.Conn.Write([]byte("\n"))
		if err != nil {
			log.Error(err)
			return
		}
	}
}