<commit_before>package node

import (
	"crypto/subtle"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"time"

	"github.com/abrander/gansoi/logger"
	"github.com/abrander/gansoi/stats"
)

// HTTPStream implements a raft stream for use with Golang's net/http.
type HTTPStream struct {
	closed   bool
	addr     string
	accepted chan net.Conn
	dial     net.Dialer
	secret   string
}

func init() {
	stats.CounterInit("http_dialed")
	stats.CounterInit("http_failed")
	stats.CounterInit("http_served")
	stats.CounterInit("http_accepted")
}

// NewHTTPStream will instantiate a new HTTPStream.
func NewHTTPStream(addr string, secret string) (*HTTPStream, error) {
	// Set up our own dialer.
	dial := net.Dialer{
		// Some crappy NAT devices will close a connection after 30 seconds of
		// inactivity. We try to keep alive every 25th second to counter this.
		KeepAlive: time.Second * 25,
	}

	h := &HTTPStream{
		addr:     addr,
		accepted: make(chan net.Conn),
		dial:     dial,
		secret:   secret,
	}

	return h, nil
}

// Dial will dial a remote http endpoint (and implement raft.StreamLayer).
func (h *HTTPStream) Dial(address string, timeout time.Duration) (net.Conn, error) {
	var conn net.Conn
	var err error
	// Make a copy of our dialer to allow custom timeout.
	dial := h.dial
	dial.Timeout = timeout

	URL, err := url.Parse(address)
	if err != nil {
		return nil, err
	}

	host, port, _ := net.SplitHostPort(URL.Host)

	stats.CounterInc("http_dialed", 1)

	logger.Green("httpstream", "Dialing %s [%s]", host, address)

	switch URL.Scheme {
	case "https":
		if port == "" {
			port = "443"
		}

		conf := &tls.Config{}
		conn, err = tls.DialWithDialer(&dial, "tcp", net.JoinHostPort(host, port), conf)

	case "http":
		if port == "" {
			port = "80"
		}

		conn, err = dial.Dial("tcp", net.JoinHostPort(host, port))

	default:
		return nil, errors.New("Unknown scheme: " + URL.Scheme)
	}

	if err != nil {
		stats.CounterInc("http_failed", 1)
		return nil, err
	}

	// We use Upgrade, and hope that will make proxies happy.
	open := fmt.Sprintf("GET /raft HTTP/1.1\nHost: %s\nUpgrade: raft-0\nSecret: %s\n\n", host, h.secret)

	_, err = conn.Write([]byte(open))
	if err != nil {
		conn.Close()

		stats.CounterInc("http_failed", 1)
		return nil, err
	}

	return conn, nil
}

// Accept waits for and returns the next connection to the listener.
func (h *HTTPStream) Accept() (net.Conn, error) {
	if h.closed {
		return nil, errors.New("Server is shutting down")
	}

	stats.CounterInc("http_accepted", 1)

	return <-h.accepted, nil
}

// Close closes the listener.
// Any blocked Accept operations will be unblocked and return errors.
func (h *HTTPStream) Close() error {
	h.closed = true

	return nil
}

// Addr returns the listener's network address.
func (h *HTTPStream) Addr() net.Addr {
	return h
}

// String implements net.Addr.
func (h *HTTPStream) String() string {
	return h.addr
}

// Network implements net.Addr.
func (h *HTTPStream) Network() string {
	return "tcp"
}

// ServeHTTP implements the http.Handler interface.
func (h *HTTPStream) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	stats.CounterInc("http_served", 1)

	if h.closed {
		http.Error(w, "Server is shutting down", http.StatusServiceUnavailable)
		return
	}

	upgrade := r.Header.Get("Upgrade")
	if upgrade != "raft-0" {
		http.Error(w, "This endpoint is for streaming raft", http.StatusBadRequest)
		return
	}

	secret := r.Header.Get("Secret")
	if subtle.ConstantTimeCompare([]byte(h.secret), []byte(secret)) != 1 {
		http.Error(w, "Wrong secret", http.StatusForbidden)
		return
	}

	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "webserver doesn't support hijacking", http.StatusInternalServerError)
		return
	}

	// We hijack the connection.
	conn, _, err := hj.Hijack()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Undo deadlines set by webserver.
	conn.SetDeadline(time.Time{})
	conn.SetWriteDeadline(time.Time{})

	h.accepted <- conn
}
<commit_msg>Support endpoints with no port specified.<commit_after>package node

import (
	"crypto/subtle"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/abrander/gansoi/logger"
	"github.com/abrander/gansoi/stats"
)

// HTTPStream implements a raft stream for use with Golang's net/http.
type HTTPStream struct {
	closed   bool
	addr     string
	accepted chan net.Conn
	dial     net.Dialer
	secret   string
}

func init() {
	stats.CounterInit("http_dialed")
	stats.CounterInit("http_failed")
	stats.CounterInit("http_served")
	stats.CounterInit("http_accepted")
}

// NewHTTPStream will instantiate a new HTTPStream.
func NewHTTPStream(addr string, secret string) (*HTTPStream, error) {
	// Set up our own dialer.
	dial := net.Dialer{
		// Some crappy NAT devices will close a connection after 30 seconds of
		// inactivity. We try to keep alive every 25th second to counter this.
		KeepAlive: time.Second * 25,
	}

	h := &HTTPStream{
		addr:     addr,
		accepted: make(chan net.Conn),
		dial:     dial,
		secret:   secret,
	}

	return h, nil
}

// Dial will dial a remote http endpoint (and implement raft.StreamLayer).
func (h *HTTPStream) Dial(address string, timeout time.Duration) (net.Conn, error) {
	var conn net.Conn
	var err error
	var host, port string
	// Make a copy of our dialer to allow custom timeout.
	dial := h.dial
	dial.Timeout = timeout

	URL, err := url.Parse(address)
	if err != nil {
		return nil, err
	}

	if strings.IndexRune(URL.Host, ':') < 0 {
		host = URL.Host
	} else {
		host, port, _ = net.SplitHostPort(URL.Host)
	}

	stats.CounterInc("http_dialed", 1)

	logger.Green("httpstream", "Dialing %s [%s]", host, address)

	switch URL.Scheme {
	case "https":
		if port == "" {
			port = "443"
		}

		conf := &tls.Config{}
		conn, err = tls.DialWithDialer(&dial, "tcp", net.JoinHostPort(host, port), conf)

	case "http":
		if port == "" {
			port = "80"
		}

		conn, err = dial.Dial("tcp", net.JoinHostPort(host, port))

	default:
		return nil, errors.New("Unknown scheme: " + URL.Scheme)
	}

	if err != nil {
		stats.CounterInc("http_failed", 1)
		return nil, err
	}

	// We use Upgrade, and hope that will make proxies happy.
	open := fmt.Sprintf("GET /raft HTTP/1.1\nHost: %s\nUpgrade: raft-0\nSecret: %s\n\n", host, h.secret)

	_, err = conn.Write([]byte(open))
	if err != nil {
		conn.Close()

		stats.CounterInc("http_failed", 1)
		return nil, err
	}

	return conn, nil
}

// Accept waits for and returns the next connection to the listener.
func (h *HTTPStream) Accept() (net.Conn, error) {
	if h.closed {
		return nil, errors.New("Server is shutting down")
	}

	stats.CounterInc("http_accepted", 1)

	return <-h.accepted, nil
}

// Close closes the listener.
// Any blocked Accept operations will be unblocked and return errors.
func (h *HTTPStream) Close() error {
	h.closed = true

	return nil
}

// Addr returns the listener's network address.
func (h *HTTPStream) Addr() net.Addr {
	return h
}

// String implements net.Addr.
func (h *HTTPStream) String() string {
	return h.addr
}

// Network implements net.Addr.
func (h *HTTPStream) Network() string {
	return "tcp"
}

// ServeHTTP implements the http.Handler interface.
func (h *HTTPStream) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	stats.CounterInc("http_served", 1)

	if h.closed {
		http.Error(w, "Server is shutting down", http.StatusServiceUnavailable)
		return
	}

	upgrade := r.Header.Get("Upgrade")
	if upgrade != "raft-0" {
		http.Error(w, "This endpoint is for streaming raft", http.StatusBadRequest)
		return
	}

	secret := r.Header.Get("Secret")
	if subtle.ConstantTimeCompare([]byte(h.secret), []byte(secret)) != 1 {
		http.Error(w, "Wrong secret", http.StatusForbidden)
		return
	}

	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "webserver doesn't support hijacking", http.StatusInternalServerError)
		return
	}

	// We hijack the connection.
	conn, _, err := hj.Hijack()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Undo deadlines set by webserver.
	conn.SetDeadline(time.Time{})
	conn.SetWriteDeadline(time.Time{})

	h.accepted <- conn
}
<|endoftext|>


<commit_before>package api

import (
	"net/http"
	"strings"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/julienschmidt/httprouter"
)

// buildHttpRoutes sets up and returns an * httprouter.Router.
// it connected the Router to the given api using the required
// parameters: requiredUserAgent and requiredPassword
func (api *API) buildHttpRoutes(requiredUserAgent string, requiredPassword string) {
	router := httprouter.New()

	router.NotFound = http.HandlerFunc(UnrecognizedCallHandler)
	router.RedirectTrailingSlash = false

	// Consensus API Calls
	if api.cs != nil {
		router.GET("/consensus", api.consensusHandler)
		router.POST("/consensus/validate/transactionset", api.consensusValidateTransactionsetHandler)
	}

	// Explorer API Calls
	if api.explorer != nil {
		router.GET("/explorer", api.explorerHandler)
		router.GET("/explorer/blocks/:height", api.explorerBlocksHandler)
		router.GET("/explorer/hashes/:hash", api.explorerHashHandler)
	}

	// Gateway API Calls
	if api.gateway != nil {
		router.GET("/gateway", api.gatewayHandler)
		router.POST("/gateway/connect/:netaddress", RequirePassword(api.gatewayConnectHandler, requiredPassword))
		router.POST("/gateway/disconnect/:netaddress", RequirePassword(api.gatewayDisconnectHandler, requiredPassword))
	}

	// Host API Calls
	if api.host != nil {
		// Calls directly pertaining to the host.
		router.GET("/host", api.hostHandlerGET)                                                   // Get the host status.
		router.POST("/host", RequirePassword(api.hostHandlerPOST, requiredPassword))              // Change the settings of the host.
		router.POST("/host/announce", RequirePassword(api.hostAnnounceHandler, requiredPassword)) // Announce the host to the network.
		router.GET("/host/estimatescore", api.hostEstimateScoreGET)

		// Calls pertaining to the storage manager that the host uses.
		router.GET("/host/storage", api.storageHandler)
		router.POST("/host/storage/folders/add", RequirePassword(api.storageFoldersAddHandler, requiredPassword))
		router.POST("/host/storage/folders/remove", RequirePassword(api.storageFoldersRemoveHandler, requiredPassword))
		router.POST("/host/storage/folders/resize", RequirePassword(api.storageFoldersResizeHandler, requiredPassword))
		router.POST("/host/storage/sectors/delete/:merkleroot", RequirePassword(api.storageSectorsDeleteHandler, requiredPassword))
	}

	// Miner API Calls
	if api.miner != nil {
		router.GET("/miner", api.minerHandler)
		router.GET("/miner/header", RequirePassword(api.minerHeaderHandlerGET, requiredPassword))
		router.POST("/miner/header", RequirePassword(api.minerHeaderHandlerPOST, requiredPassword))
		router.GET("/miner/start", RequirePassword(api.minerStartHandler, requiredPassword))
		router.GET("/miner/stop", RequirePassword(api.minerStopHandler, requiredPassword))
	}

	// Renter API Calls
	if api.renter != nil {
		router.GET("/renter", api.renterHandlerGET)
		router.POST("/renter", RequirePassword(api.renterHandlerPOST, requiredPassword))
		router.GET("/renter/contracts", api.renterContractsHandler)
		router.GET("/renter/downloads", api.renterDownloadsHandler)
		router.GET("/renter/files", api.renterFilesHandler)
		router.GET("/renter/prices", api.renterPricesHandler)

		// TODO: re-enable these routes once the new .sia format has been
		// standardized and implemented.
		// router.POST("/renter/load", RequirePassword(api.renterLoadHandler, requiredPassword))
		// router.POST("/renter/loadascii", RequirePassword(api.renterLoadAsciiHandler, requiredPassword))
		// router.GET("/renter/share", RequirePassword(api.renterShareHandler, requiredPassword))
		// router.GET("/renter/shareascii", RequirePassword(api.renterShareAsciiHandler, requiredPassword))

		router.POST("/renter/delete/*siapath", RequirePassword(api.renterDeleteHandler, requiredPassword))
		router.GET("/renter/download/*siapath", RequirePassword(api.renterDownloadHandler, requiredPassword))
		router.GET("/renter/downloadasync/*siapath", RequirePassword(api.renterDownloadAsyncHandler, requiredPassword))
		router.POST("/renter/rename/*siapath", RequirePassword(api.renterRenameHandler, requiredPassword))
		router.POST("/renter/upload/*siapath", RequirePassword(api.renterUploadHandler, requiredPassword))

		// HostDB endpoints.
		router.GET("/hostdb/active", api.hostdbActiveHandler)
		router.GET("/hostdb/all", api.hostdbAllHandler)
		router.GET("/hostdb/hosts/:pubkey", api.hostdbHostsHandler)
	}

	// Transaction pool API Calls
	if api.tpool != nil {
		router.GET("/tpool/fee", api.tpoolFeeHandlerGET)
		router.GET("/tpool/raw/:id", api.tpoolRawHandlerGET)
		router.POST("/tpool/raw", api.tpoolRawHandlerPOST)
		router.GET("/tpool/confirmed/:id", api.tpoolConfirmedGET)

		// TODO: re-enable this route once the transaction pool API has been finalized
		//router.GET("/transactionpool/transactions", api.transactionpoolTransactionsHandler)
	}

	// Wallet API Calls
	if api.wallet != nil {
		router.GET("/wallet", api.walletHandler)
		router.POST("/wallet/033x", RequirePassword(api.wallet033xHandler, requiredPassword))
		router.GET("/wallet/address", RequirePassword(api.walletAddressHandler, requiredPassword))
		router.GET("/wallet/addresses", api.walletAddressesHandler)
		router.GET("/wallet/backup", RequirePassword(api.walletBackupHandler, requiredPassword))
		router.POST("/wallet/init", RequirePassword(api.walletInitHandler, requiredPassword))
		router.POST("/wallet/init/seed", RequirePassword(api.walletInitSeedHandler, requiredPassword))
		router.POST("/wallet/lock", RequirePassword(api.walletLockHandler, requiredPassword))
		router.POST("/wallet/seed", RequirePassword(api.walletSeedHandler, requiredPassword))
		router.GET("/wallet/seeds", RequirePassword(api.walletSeedsHandler, requiredPassword))
		router.POST("/wallet/siacoins", RequirePassword(api.walletSiacoinsHandler, requiredPassword))
		router.POST("/wallet/siafunds", RequirePassword(api.walletSiafundsHandler, requiredPassword))
		router.POST("/wallet/siagkey", RequirePassword(api.walletSiagkeyHandler, requiredPassword))
		router.POST("/wallet/sweep/seed", RequirePassword(api.walletSweepSeedHandler, requiredPassword))
		router.GET("/wallet/transaction/:id", api.walletTransactionHandler)
		router.GET("/wallet/transactions", api.walletTransactionsHandler)
		router.GET("/wallet/transactions/:addr", api.walletTransactionsAddrHandler)
		router.GET("/wallet/verify/address/:addr", api.walletVerifyAddressHandler)
		router.POST("/wallet/unlock", RequirePassword(api.walletUnlockHandler, requiredPassword))
		router.POST("/wallet/changepassword", RequirePassword(api.walletChangePasswordHandler, requiredPassword))
	}

	// Apply UserAgent middleware and return the Router
	api.router = cleanCloseHandler(RequireUserAgent(router, requiredUserAgent))
	return
}

// cleanCloseHandler wraps the entire API, ensuring that underlying conns are
// not leaked if the remote end closes the connection before the underlying
// handler finishes.
func cleanCloseHandler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Close this file handle either when the function completes or when the
		// connection is done.
		done := make(chan struct{})
		go func(w http.ResponseWriter, r *http.Request) {
			defer close(done)
			next.ServeHTTP(w, r)
		}(w, r)
		select {
		case <-done:
		case <-r.Context().Done():
		}

		// Sanity check - thread should not take more than an hour to return. This
		// must be done in a goroutine, otherwise the server will not close the
		// underlying socket for this API call.
		timer := time.NewTimer(time.Minute * 60)
		go func() {
			select {
			case <-done:
				timer.Stop()
			case <-timer.C:
				build.Severe("api call is taking more than 60 minutes to return:", r.URL.Path)
			}
		}()
	})
}

// RequireUserAgent is middleware that requires all requests to set a
// UserAgent that contains the specified string.
func RequireUserAgent(h http.Handler, ua string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		if !strings.Contains(req.UserAgent(), ua) {
			WriteError(w, Error{"Browser access disabled due to security vulnerability. Use Sia-UI or siac."}, http.StatusBadRequest)
			return
		}
		h.ServeHTTP(w, req)
	})
}

// RequirePassword is middleware that requires a request to authenticate with a
// password using HTTP basic auth. Usernames are ignored. Empty passwords
// indicate no authentication is required.
func RequirePassword(h httprouter.Handle, password string) httprouter.Handle {
	// An empty password is equivalent to no password.
	if password == "" {
		return h
	}
	return func(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
		_, pass, ok := req.BasicAuth()
		if !ok || pass != password {
			w.Header().Set("WWW-Authenticate", "Basic realm=\"SiaAPI\"")
			WriteError(w, Error{"API authentication failed."}, http.StatusUnauthorized)
			return
		}
		h(w, req, ps)
	}
}
<commit_msg>fix http server race<commit_after>package api

import (
	"net/http"
	"strings"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/julienschmidt/httprouter"
)

// buildHttpRoutes sets up and returns an * httprouter.Router.
// it connected the Router to the given api using the required
// parameters: requiredUserAgent and requiredPassword
func (api *API) buildHttpRoutes(requiredUserAgent string, requiredPassword string) {
	router := httprouter.New()

	router.NotFound = http.HandlerFunc(UnrecognizedCallHandler)
	router.RedirectTrailingSlash = false

	// Consensus API Calls
	if api.cs != nil {
		router.GET("/consensus", api.consensusHandler)
		router.POST("/consensus/validate/transactionset", api.consensusValidateTransactionsetHandler)
	}

	// Explorer API Calls
	if api.explorer != nil {
		router.GET("/explorer", api.explorerHandler)
		router.GET("/explorer/blocks/:height", api.explorerBlocksHandler)
		router.GET("/explorer/hashes/:hash", api.explorerHashHandler)
	}

	// Gateway API Calls
	if api.gateway != nil {
		router.GET("/gateway", api.gatewayHandler)
		router.POST("/gateway/connect/:netaddress", RequirePassword(api.gatewayConnectHandler, requiredPassword))
		router.POST("/gateway/disconnect/:netaddress", RequirePassword(api.gatewayDisconnectHandler, requiredPassword))
	}

	// Host API Calls
	if api.host != nil {
		// Calls directly pertaining to the host.
		router.GET("/host", api.hostHandlerGET)                                                   // Get the host status.
		router.POST("/host", RequirePassword(api.hostHandlerPOST, requiredPassword))              // Change the settings of the host.
		router.POST("/host/announce", RequirePassword(api.hostAnnounceHandler, requiredPassword)) // Announce the host to the network.
		router.GET("/host/estimatescore", api.hostEstimateScoreGET)

		// Calls pertaining to the storage manager that the host uses.
		router.GET("/host/storage", api.storageHandler)
		router.POST("/host/storage/folders/add", RequirePassword(api.storageFoldersAddHandler, requiredPassword))
		router.POST("/host/storage/folders/remove", RequirePassword(api.storageFoldersRemoveHandler, requiredPassword))
		router.POST("/host/storage/folders/resize", RequirePassword(api.storageFoldersResizeHandler, requiredPassword))
		router.POST("/host/storage/sectors/delete/:merkleroot", RequirePassword(api.storageSectorsDeleteHandler, requiredPassword))
	}

	// Miner API Calls
	if api.miner != nil {
		router.GET("/miner", api.minerHandler)
		router.GET("/miner/header", RequirePassword(api.minerHeaderHandlerGET, requiredPassword))
		router.POST("/miner/header", RequirePassword(api.minerHeaderHandlerPOST, requiredPassword))
		router.GET("/miner/start", RequirePassword(api.minerStartHandler, requiredPassword))
		router.GET("/miner/stop", RequirePassword(api.minerStopHandler, requiredPassword))
	}

	// Renter API Calls
	if api.renter != nil {
		router.GET("/renter", api.renterHandlerGET)
		router.POST("/renter", RequirePassword(api.renterHandlerPOST, requiredPassword))
		router.GET("/renter/contracts", api.renterContractsHandler)
		router.GET("/renter/downloads", api.renterDownloadsHandler)
		router.GET("/renter/files", api.renterFilesHandler)
		router.GET("/renter/prices", api.renterPricesHandler)

		// TODO: re-enable these routes once the new .sia format has been
		// standardized and implemented.
		// router.POST("/renter/load", RequirePassword(api.renterLoadHandler, requiredPassword))
		// router.POST("/renter/loadascii", RequirePassword(api.renterLoadAsciiHandler, requiredPassword))
		// router.GET("/renter/share", RequirePassword(api.renterShareHandler, requiredPassword))
		// router.GET("/renter/shareascii", RequirePassword(api.renterShareAsciiHandler, requiredPassword))

		router.POST("/renter/delete/*siapath", RequirePassword(api.renterDeleteHandler, requiredPassword))
		router.GET("/renter/download/*siapath", RequirePassword(api.renterDownloadHandler, requiredPassword))
		router.GET("/renter/downloadasync/*siapath", RequirePassword(api.renterDownloadAsyncHandler, requiredPassword))
		router.POST("/renter/rename/*siapath", RequirePassword(api.renterRenameHandler, requiredPassword))
		router.POST("/renter/upload/*siapath", RequirePassword(api.renterUploadHandler, requiredPassword))

		// HostDB endpoints.
		router.GET("/hostdb/active", api.hostdbActiveHandler)
		router.GET("/hostdb/all", api.hostdbAllHandler)
		router.GET("/hostdb/hosts/:pubkey", api.hostdbHostsHandler)
	}

	// Transaction pool API Calls
	if api.tpool != nil {
		router.GET("/tpool/fee", api.tpoolFeeHandlerGET)
		router.GET("/tpool/raw/:id", api.tpoolRawHandlerGET)
		router.POST("/tpool/raw", api.tpoolRawHandlerPOST)
		router.GET("/tpool/confirmed/:id", api.tpoolConfirmedGET)

		// TODO: re-enable this route once the transaction pool API has been finalized
		//router.GET("/transactionpool/transactions", api.transactionpoolTransactionsHandler)
	}

	// Wallet API Calls
	if api.wallet != nil {
		router.GET("/wallet", api.walletHandler)
		router.POST("/wallet/033x", RequirePassword(api.wallet033xHandler, requiredPassword))
		router.GET("/wallet/address", RequirePassword(api.walletAddressHandler, requiredPassword))
		router.GET("/wallet/addresses", api.walletAddressesHandler)
		router.GET("/wallet/backup", RequirePassword(api.walletBackupHandler, requiredPassword))
		router.POST("/wallet/init", RequirePassword(api.walletInitHandler, requiredPassword))
		router.POST("/wallet/init/seed", RequirePassword(api.walletInitSeedHandler, requiredPassword))
		router.POST("/wallet/lock", RequirePassword(api.walletLockHandler, requiredPassword))
		router.POST("/wallet/seed", RequirePassword(api.walletSeedHandler, requiredPassword))
		router.GET("/wallet/seeds", RequirePassword(api.walletSeedsHandler, requiredPassword))
		router.POST("/wallet/siacoins", RequirePassword(api.walletSiacoinsHandler, requiredPassword))
		router.POST("/wallet/siafunds", RequirePassword(api.walletSiafundsHandler, requiredPassword))
		router.POST("/wallet/siagkey", RequirePassword(api.walletSiagkeyHandler, requiredPassword))
		router.POST("/wallet/sweep/seed", RequirePassword(api.walletSweepSeedHandler, requiredPassword))
		router.GET("/wallet/transaction/:id", api.walletTransactionHandler)
		router.GET("/wallet/transactions", api.walletTransactionsHandler)
		router.GET("/wallet/transactions/:addr", api.walletTransactionsAddrHandler)
		router.GET("/wallet/verify/address/:addr", api.walletVerifyAddressHandler)
		router.POST("/wallet/unlock", RequirePassword(api.walletUnlockHandler, requiredPassword))
		router.POST("/wallet/changepassword", RequirePassword(api.walletChangePasswordHandler, requiredPassword))
	}

	// Apply UserAgent middleware and return the Router
	api.router = cleanCloseHandler(RequireUserAgent(router, requiredUserAgent))
	return
}

// cleanCloseHandler wraps the entire API, ensuring that underlying conns are
// not leaked if the remote end closes the connection before the underlying
// handler finishes.
func cleanCloseHandler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Close this file handle either when the function completes or when the
		// connection is done.
		done := make(chan struct{})
		go func(w http.ResponseWriter, r *http.Request) {
			defer close(done)
			next.ServeHTTP(w, r)
		}(w, r)
		select {
		case <-done:
		}

		// Sanity check - thread should not take more than an hour to return. This
		// must be done in a goroutine, otherwise the server will not close the
		// underlying socket for this API call.
		timer := time.NewTimer(time.Minute * 60)
		go func() {
			select {
			case <-done:
				timer.Stop()
			case <-timer.C:
				build.Severe("api call is taking more than 60 minutes to return:", r.URL.Path)
			}
		}()
	})
}

// RequireUserAgent is middleware that requires all requests to set a
// UserAgent that contains the specified string.
func RequireUserAgent(h http.Handler, ua string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		if !strings.Contains(req.UserAgent(), ua) {
			WriteError(w, Error{"Browser access disabled due to security vulnerability. Use Sia-UI or siac."}, http.StatusBadRequest)
			return
		}
		h.ServeHTTP(w, req)
	})
}

// RequirePassword is middleware that requires a request to authenticate with a
// password using HTTP basic auth. Usernames are ignored. Empty passwords
// indicate no authentication is required.
func RequirePassword(h httprouter.Handle, password string) httprouter.Handle {
	// An empty password is equivalent to no password.
	if password == "" {
		return h
	}
	return func(w http.ResponseWriter, req *http.Request, ps httprouter.Params) {
		_, pass, ok := req.BasicAuth()
		if !ok || pass != password {
			w.Header().Set("WWW-Authenticate", "Basic realm=\"SiaAPI\"")
			WriteError(w, Error{"API authentication failed."}, http.StatusUnauthorized)
			return
		}
		h(w, req, ps)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Nodetemple <hostmaster@nodetemple.com>

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"fmt"
	"strings"
	"text/template"

	flag "github.com/ogier/pflag"
)

var (
	cmdHelp = &Command{
		Name:        "help",
		Usage:       "<command>",
		Summary:     "Show a list of commands or help for one command",
		Description: "Show a list of commands or detailed help for one command",
		Run:         runHelp,
	}

	globalUsageTemplate  *template.Template
	commandUsageTemplate *template.Template
	templFuncs           = template.FuncMap{
		"descToLines": func(s string) []string {
			return strings.Split(strings.Trim(s, "\n\t "), "\n")
		},
		"printFlag": func(shorthand, name, defvalue, usage string) string {
			format := "--%s=%s\t%s"
			/*if _, ok := flag.Value.(*stringValue); ok {
				format = "--%s=%q\t%s"
			}*/
			if len(shorthand) > 0 {
				format = "    -%s, " + format
			} else {
				format = "     %s   " + format
			}
			return fmt.Sprintf(format, shorthand, name, defvalue, usage)
		},
	}
)

func init() {
	globalUsageTemplate = template.Must(template.New("global_usage").Funcs(templFuncs).Parse(`
NAME:
{{printf "\t%s - %s" .Executable .Description}}

USAGE:
{{printf "\t%s" .Executable}} [global flags] <command> [command flags] [arguments...]

VERSION:
{{printf "\t%s" .Version}}

COMMANDS:{{range .Commands}}
{{printf "\t%s\t%s" .Name .Summary}}{{end}}

GLOBAL FLAGS:{{range .Flags}}
{{printFlag .Shorthand .Name .DefValue .Usage}}{{end}}

Global flags can also be configured via upper-case environment variables prefixed with "{{.ExeEnvPrefix}}_"
For example, "--some-flag" => "{{.ExeEnvPrefix}}_SOME_FLAG"

Run "{{.Executable}} help <command>" for more details on a specific command
`[1:]))
	commandUsageTemplate = template.Must(template.New("command_usage").Funcs(templFuncs).Parse(`
NAME:
{{printf "\t%s - %s" .Cmd.Name .Cmd.Summary}}

USAGE:
{{printf "\t%s %s %s" .Executable .Cmd.Name .Cmd.Usage}}

{{if .Cmd.Description}}DESCRIPTION:{{range $line := descToLines .Cmd.Description}}
{{printf "\t%s" $line}}{{end}}

{{end}}{{if .Cmd.Subcommands}}COMMANDS:{{range .Cmd.Subcommands}}
{{printf "\t%s\t%s" .Name .Summary}}{{end}}

{{end}}{{if .CmdFlags}}FLAGS:{{range .CmdFlags}}
{{printFlag .Shorthand .Name .DefValue .Usage}}{{end}}

{{end}}For help on global flags run "{{.Executable}} help"
`[1:]))
}

func runHelp(args []string) int {
	if len(args) < 1 {
		printGlobalUsage()
		return OK
	}

	var cmd *Command

	for _, c := range commands {
		if c.Name == args[0] {
			cmd = c
			break
		}
	}

	if cmd == nil {
		stderr("Error: unknown command: %s\n", args[0])
		stderr("Run '%s help' for usage information\n", cliName)
		return ERROR_NO_COMMAND
	}

	printCommandUsage(cmd)
	return OK
}

func printGlobalUsage() {
	globalUsageTemplate.Execute(out, struct {
		Executable  string
		Commands    []*Command
		Flags       []*flag.Flag
		Description string
		Version     string
	}{
		cliName,
		commands,
		getAllFlags(),
		cliDescription,
		cliVersion,
	})
}

func printCommandUsage(cmd *Command) {
	commandUsageTemplate.Execute(out, struct {
		Executable   string
		ExeEnvPrefix string
		Cmd          *Command
		CmdFlags     []*flag.Flag
	}{
		cliName,
		strings.ToUpper(cliName),
		cmd,
		getFlags(&cmd.Flags),
	})
}

func getAllFlags() (flags []*flag.Flag) {
	return getFlags(globalFlagSet)
}

func getFlags(flagset *flag.FlagSet) (flags []*flag.Flag) {
	flags = make([]*flag.Flag, 0)
	flagset.VisitAll(func(f *flag.Flag) {
		flags = append(flags, f)
	})
	return
}
<commit_msg>Fix ExeEnvPrefix<commit_after>/*
Copyright 2015 Nodetemple <hostmaster@nodetemple.com>

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"fmt"
	"strings"
	"text/template"

	flag "github.com/ogier/pflag"
)

var (
	cmdHelp = &Command{
		Name:        "help",
		Usage:       "<command>",
		Summary:     "Show a list of commands or help for one command",
		Description: "Show a list of commands or detailed help for one command",
		Run:         runHelp,
	}

	globalUsageTemplate  *template.Template
	commandUsageTemplate *template.Template
	templFuncs           = template.FuncMap{
		"descToLines": func(s string) []string {
			return strings.Split(strings.Trim(s, "\n\t "), "\n")
		},
		"printFlag": func(shorthand, name, defvalue, usage string) string {
			format := "--%s=%s\t%s"
			/*if _, ok := flag.Value.(*stringValue); ok {
				format = "--%s=%q\t%s"
			}*/
			if len(shorthand) > 0 {
				format = "    -%s, " + format
			} else {
				format = "     %s   " + format
			}
			return fmt.Sprintf(format, shorthand, name, defvalue, usage)
		},
	}
)

func init() {
	globalUsageTemplate = template.Must(template.New("global_usage").Funcs(templFuncs).Parse(`
NAME:
{{printf "\t%s - %s" .Executable .Description}}

USAGE:
{{printf "\t%s" .Executable}} [global flags] <command> [command flags] [arguments...]

VERSION:
{{printf "\t%s" .Version}}

COMMANDS:{{range .Commands}}
{{printf "\t%s\t%s" .Name .Summary}}{{end}}

GLOBAL FLAGS:{{range .Flags}}
{{printFlag .Shorthand .Name .DefValue .Usage}}{{end}}

Global flags can also be configured via upper-case environment variables prefixed with "{{.ExeEnvPrefix}}_"
For example: "--some-flag" => "{{.ExeEnvPrefix}}_SOME_FLAG"

Run "{{.Executable}} help <command>" for more details on a specific command
`[1:]))
	commandUsageTemplate = template.Must(template.New("command_usage").Funcs(templFuncs).Parse(`
NAME:
{{printf "\t%s - %s" .Cmd.Name .Cmd.Summary}}

USAGE:
{{printf "\t%s %s %s" .Executable .Cmd.Name .Cmd.Usage}}

{{if .Cmd.Description}}DESCRIPTION:{{range $line := descToLines .Cmd.Description}}
{{printf "\t%s" $line}}{{end}}

{{end}}{{if .Cmd.Subcommands}}COMMANDS:{{range .Cmd.Subcommands}}
{{printf "\t%s\t%s" .Name .Summary}}{{end}}

{{end}}{{if .CmdFlags}}FLAGS:{{range .CmdFlags}}
{{printFlag .Shorthand .Name .DefValue .Usage}}{{end}}

{{end}}For help on global flags run "{{.Executable}} help"
`[1:]))
}

func runHelp(args []string) int {
	if len(args) < 1 {
		printGlobalUsage()
		return OK
	}

	var cmd *Command

	for _, c := range commands {
		if c.Name == args[0] {
			cmd = c
			break
		}
	}

	if cmd == nil {
		stderr("Error: unknown command: %s\n", args[0])
		stderr("Run '%s help' for usage information\n", cliName)
		return ERROR_NO_COMMAND
	}

	printCommandUsage(cmd)
	return OK
}

func printGlobalUsage() {
	globalUsageTemplate.Execute(out, struct {
		Executable  string
		ExeEnvPrefix string
		Commands    []*Command
		Flags       []*flag.Flag
		Description string
		Version     string
	}{
		cliName,
		strings.ToUpper(cliName),
		commands,
		getAllFlags(),
		cliDescription,
		cliVersion,
	})
}

func printCommandUsage(cmd *Command) {
	commandUsageTemplate.Execute(out, struct {
		Executable   string
		Cmd          *Command
		CmdFlags     []*flag.Flag
	}{
		cliName,
		cmd,
		getFlags(&cmd.Flags),
	})
}

func getAllFlags() (flags []*flag.Flag) {
	return getFlags(globalFlagSet)
}

func getFlags(flagset *flag.FlagSet) (flags []*flag.Flag) {
	flags = make([]*flag.Flag, 0)
	flagset.VisitAll(func(f *flag.Flag) {
		flags = append(flags, f)
	})
	return
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/aidarkhanov/nanoid"
	"github.com/jackc/pgx/v4"
	"github.com/jackc/pgx/v4/pgxpool"
	"github.com/rotisserie/eris"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"

	"github.com/ngld/knossos/packages/server/pkg/config"
	"github.com/ngld/knossos/packages/server/pkg/db"
	"github.com/ngld/knossos/packages/server/pkg/db/queries"
	"github.com/ngld/knossos/packages/server/pkg/importer"
	"github.com/ngld/knossos/packages/server/pkg/nblog"
)

func handleFile(ctx context.Context, q *queries.DBQuerier, url string) int32 {
	if url == "" {
		return 0
	}

	fid, err := q.CreateExternalFile(ctx, queries.CreateExternalFileParams{
		StorageKey: "ext#" + nanoid.New(),
		Filesize:   0,
		Public:     true,
		External:   []string{url},
		Owner:      1,
	})
	if err != nil {
		log.Fatal().Err(err).Msgf("Failed to store external file %s.", url)
	}

	return fid
}

type FilesSource struct {
	packageID int32
	files     []importer.KnFile
	curIdx    int
	length    int
	err       error
}

func (fs *FilesSource) Next() bool {
	fs.curIdx++
	return fs.err == nil && fs.curIdx < fs.length
}

var filesSourceCols = []string{"package_id", "path", "archive", "archive_path", "checksum_algo", "checksum_digest"}

func (fs *FilesSource) Values() ([]interface{}, error) {
	file := fs.files[fs.curIdx]

	digest, err := hex.DecodeString(file.Checksum[1])
	if err != nil {
		fs.err = err
		return nil, err
	}

	result := make([]interface{}, 6)
	result[0] = fs.packageID
	result[1] = file.Filename
	result[2] = file.Archive
	result[3] = file.OrigName
	result[4] = file.Checksum[0]
	result[5] = digest

	return result, nil
}

func (fs *FilesSource) Err() error {
	return fs.err
}

func CopyFromFiles(ctx context.Context, tx pgx.Tx, pkg int32, files []importer.KnFile) (int64, error) {
	source := FilesSource{
		packageID: pkg,
		files:     files,
		curIdx:    -1,
		length:    len(files),
		err:       nil,
	}

	return tx.CopyFrom(ctx, pgx.Identifier{"mod_package_files"}, filesSourceCols, &source)
}

func main() {
	cfg, loader := config.Loader()
	// flags := loader.Flags()

	if err := loader.Load(); err != nil {
		if strings.Contains(err.Error(), "help requested") {
			os.Exit(3)
		}

		panic(err)
	}

	writer := zerolog.ConsoleWriter{Out: os.Stdout}
	writer.TimeFormat = "02.01.2006 15:04:05 MST"
	writer.FormatFieldValue = func(value interface{}) string {
		if value == nil {
			return "                "
		}

		str, ok := value.(string)
		if ok {
			if len(str) == 16 {
				// color request IDs in cyan  we have to guess based on the field content because we can't get
				// the current field name
				return fmt.Sprintf("\x1b[%dm%s\x1b[0m", 36, value)
			} else if strings.Contains(str, "\\n") && strings.Contains(str, "\\t") {
				// unquote values that contain line breaks and tabs because they're most likely stack traces
				str, err := strconv.Unquote(str)
				if err == nil {
					return str
				}
			}
		}

		return fmt.Sprintf("%s", value)
	}

	log.Logger = log.Output(writer)

	zerolog.ErrorStackMarshaler = func(err error) interface{} {
		return eris.ToString(err, true)
	}

	if err := cfg.Validate(); err != nil {
		log.Fatal().Err(err).Msg("Failed to parse config")
	}

	zerolog.SetGlobalLevel(cfg.LogLevel())
	log.Logger = log.Logger.With().Caller().Stack().Logger()

	dbConfig, err := pgxpool.ParseConfig(cfg.Database)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to parse DB config")
	}

	dbConfig.ConnConfig.Logger = nblog.PgxLogger{}
	dbConfig.ConnConfig.LogLevel = pgx.LogLevelDebug
	pool, err := pgxpool.ConnectConfig(context.Background(), dbConfig)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to open DB connection")
	}

	log.Info().Msg("Finished parsing configuration; loading data")

	inFile, err := os.Open("repo.json")
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to open repo.json")
	}
	defer inFile.Close()

	jsonStr, err := ioutil.ReadAll(inFile)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to read repo.json")
	}

	data := importer.KnStruct{}
	err = json.Unmarshal(jsonStr, &data)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to parse repo.json")
	}

	typeMap := map[string]db.ModType{
		"mod":    db.TypeMod,
		"tc":     db.TypeTotalConversion,
		"engine": db.TypeEngine,
		"tool":   db.TypeTool,
		"ext":    db.TypeExtension,
	}

	stabilityMap := map[string]db.EngineStability{
		"":        db.EngineUnknown,
		"stable":  db.EngineStable,
		"rc":      db.EngineRC,
		"nightly": db.EngineNightly,
	}

	pkgTypeMap := map[string]db.PackageType{
		"required":    db.PackageRequired,
		"recommended": db.PackageRecommended,
		"optional":    db.PackageOptional,
	}

	ctx := context.Background()
	tx, err := pool.Begin(ctx)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to initiate transaction")
	}

	_, err = tx.Exec(ctx, "SET CONSTRAINTS ALL DEFERRED")
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to setup the transaction")
	}

	q := queries.NewQuerier(tx)

	count := len(data.Mods)
	aidCache := map[string]int32{}
	for i, mod := range data.Mods {
		log.Info().Msgf("Mod %4d of %d", i, count)

		aid, ok := aidCache[mod.ID]
		if !ok {
			mType, ok := typeMap[mod.Type]
			if !ok {
				log.Warn().Msgf("Invalid mod type %s detected for mod %s (%s)", mod.Type, mod.Title, mod.ID)
				mType = db.TypeMod
			}

			aid, err = q.CreateMod(ctx, queries.CreateModParams{
				Modid: mod.ID,
				Title: mod.Title,
				Type:  int16(mType),
			})
			if err != nil {
				log.Fatal().Err(err).Msgf("Failed to insert mod %s (%s)", mod.Title, mod.ID)
			}

			aidCache[mod.ID] = aid
		}

		log.Info().Msgf("AID %d", aid)

		rStabilility, ok := stabilityMap[mod.Stability]
		if !ok {
			log.Warn().Msgf("Invalid stability %s detected for mod %s (%s)", mod.Stability, mod.ID, mod.Version)
			rStabilility = db.EngineUnknown
		}

		// TODO files (logo -> ?, tile -> teaser, banner -> banner, screenshots)
		relDate, err := time.Parse("2006-01-02T15:04:05Z", mod.FirstRelease)
		if err != nil {
			if mod.FirstRelease == "" {
				relDate = time.Unix(1, 0)
			} else {
				log.Fatal().Msgf("Invalid release date for mod %s (%s): %s", mod.ID, mod.Version, mod.FirstRelease)
			}
		}

		updateDate, err := time.Parse("2006-01-02T15:04:05Z", mod.LastUpdate)
		if err != nil {
			if mod.LastUpdate == "" {
				updateDate = time.Unix(1, 0)
			} else {
				log.Fatal().Msgf("Invalid update date for mod %s (%s): %s", mod.ID, mod.Version, mod.LastUpdate)
			}
		}

		screens := make([]int32, len(mod.Screenshots))
		for idx, url := range mod.Screenshots {
			screens[idx] = handleFile(ctx, q, url)
		}

		modParams := queries.CreateReleaseParams{
			ModAid:        aid,
			Version:       mod.Version,
			Stability:     int16(rStabilility),
			Description:   mod.Description,
			ReleaseThread: mod.ReleaseThread,
			Videos:        mod.Videos,
			Notes:         mod.Notes,
			Cmdline:       mod.Cmdline,
			Screenshots:   screens,
			Teaser:        handleFile(ctx, q, mod.Tile),
			Banner:        handleFile(ctx, q, mod.Banner),
		}

		modParams.Released.Set(relDate)
		modParams.Updated.Set(updateDate)

		rid, err := q.CreateRelease(ctx, modParams)
		if err != nil {
			log.Fatal().Err(err).Msgf("Failed to insert release %s (%s)", mod.ID, mod.Version)
		}

		for _, pkg := range mod.Packages {
			pkgType, ok := pkgTypeMap[pkg.Status]
			if !ok {
				log.Fatal().Msgf("Failed to parse package status %s for %s in %s (%s)", pkg.Status, pkg.Name, mod.ID, mod.Version)
			}

			pid, err := q.CreatePackage(ctx, queries.CreatePackageParams{
				ReleaseID: rid,
				Name:      pkg.Name,
				Folder:    pkg.Folder,
				Notes:     pkg.Notes,
				Type:      int16(pkgType),
				// TODO: Actually parse the env spec
				CpuSpecs:  strings.Split(pkg.Environment, " && "),
				KnossosVp: pkg.IsVp,
			})
			if err != nil {
				log.Fatal().Err(err).Msgf("Failed to insert package %s for %s (%s)", pkg.Name, mod.ID, mod.Version)
			}

			pkgBatch := new(pgx.Batch)

			for _, dep := range pkg.Dependencies {
				q.CreatePackageDependencyBatch(pkgBatch, queries.CreatePackageDependencyParams{
					PackageID: pid,
					Modid:     dep.ID,
					Version:   dep.Version,
					Packages:  dep.Packages,
				})
			}

			for _, exe := range pkg.Executables {
				prio := int16(0)
				// See https://github.com/ngld/knossos/blob/1f60d925498c02d3db76a54d3ee20c31b75c5a21/knossos/repo.py#L35-L40
				if exe.Properties.X64 {
					prio += 50
				}
				if exe.Properties.AVX2 {
					prio += 3
				}
				if exe.Properties.AVX {
					prio += 2
				}
				if exe.Properties.SSE2 {
					prio++
				}

				q.CreatePackageExecutableBatch(pkgBatch, queries.CreatePackageExecutableParams{
					PackageID: pid,
					Path:      exe.File,
					Label:     exe.Label,
					Priority:  prio,
					Debug:     strings.Contains(exe.Label, "Debug"),
				})
			}

			for _, archive := range pkg.Files {
				digest, err := hex.DecodeString(archive.Checksum[1])
				if err != nil {
					log.Fatal().Err(err).Msgf("Failed to parse digest for archive %s of %s for %s (%s)", archive.Filename, pkg.Name, mod.ID, mod.Version)
				}

				fid, err := q.CreateExternalFile(ctx, queries.CreateExternalFileParams{
					StorageKey: "ext#" + nanoid.New(),
					Filesize:   int32(archive.FileSize),
					Public:     true,
					External:   archive.URLs,
					Owner:      1,
				})
				if err != nil {
					log.Fatal().Err(err).Msgf("Failed to index file for archive %s of %s for %s (%s)", archive.Filename, pkg.Name, mod.ID, mod.Version)
				}

				q.CreatePackageArchiveBatch(pkgBatch, queries.CreatePackageArchiveParams{
					PackageID:      pid,
					Label:          archive.Filename,
					Destination:    archive.Dest,
					ChecksumAlgo:   archive.Checksum[0],
					ChecksumDigest: digest,
					FileID:         fid,
				})
			}

			count, err := CopyFromFiles(ctx, tx, pid, pkg.Filelist)
			if err != nil {
				log.Fatal().Err(err).Msgf("Failed to process filelist on %s for %s (%s)", pkg.Name, mod.ID, mod.Version)
			}
			if count != int64(len(pkg.Filelist)) {
				log.Fatal().Err(err).Msgf("Imported only %d of %d files on %s for %s (%s)", count, len(pkg.Filelist), pkg.Name, mod.ID, mod.Version)
			}

			batchResults := tx.SendBatch(ctx, pkgBatch)
			for idx := 0; idx < pkgBatch.Len(); idx++ {
				_, err := batchResults.Exec()
				if err != nil {
					log.Fatal().Err(err).Msgf("Failed process package contents for %s of %s (%s)", pkg.Name, mod.ID, mod.Version)
				}
			}
			batchResults.Close()
		}
	}

	log.Info().Msg("Comitting")
	err = tx.Commit(ctx)
	if err != nil {
		log.Fatal().Err(err).Msg("Commit failed")
	}

	log.Info().Msg("Done")
}
<commit_msg>fix(nebula): correct the importer's date format<commit_after>package main

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/aidarkhanov/nanoid"
	"github.com/jackc/pgx/v4"
	"github.com/jackc/pgx/v4/pgxpool"
	"github.com/rotisserie/eris"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"

	"github.com/ngld/knossos/packages/server/pkg/config"
	"github.com/ngld/knossos/packages/server/pkg/db"
	"github.com/ngld/knossos/packages/server/pkg/db/queries"
	"github.com/ngld/knossos/packages/server/pkg/importer"
	"github.com/ngld/knossos/packages/server/pkg/nblog"
)

func handleFile(ctx context.Context, q *queries.DBQuerier, url string) int32 {
	if url == "" {
		return 0
	}

	fid, err := q.CreateExternalFile(ctx, queries.CreateExternalFileParams{
		StorageKey: "ext#" + nanoid.New(),
		Filesize:   0,
		Public:     true,
		External:   []string{url},
		Owner:      1,
	})
	if err != nil {
		log.Fatal().Err(err).Msgf("Failed to store external file %s.", url)
	}

	return fid
}

type FilesSource struct {
	packageID int32
	files     []importer.KnFile
	curIdx    int
	length    int
	err       error
}

func (fs *FilesSource) Next() bool {
	fs.curIdx++
	return fs.err == nil && fs.curIdx < fs.length
}

var filesSourceCols = []string{"package_id", "path", "archive", "archive_path", "checksum_algo", "checksum_digest"}

func (fs *FilesSource) Values() ([]interface{}, error) {
	file := fs.files[fs.curIdx]

	digest, err := hex.DecodeString(file.Checksum[1])
	if err != nil {
		fs.err = err
		return nil, err
	}

	result := make([]interface{}, 6)
	result[0] = fs.packageID
	result[1] = file.Filename
	result[2] = file.Archive
	result[3] = file.OrigName
	result[4] = file.Checksum[0]
	result[5] = digest

	return result, nil
}

func (fs *FilesSource) Err() error {
	return fs.err
}

func CopyFromFiles(ctx context.Context, tx pgx.Tx, pkg int32, files []importer.KnFile) (int64, error) {
	source := FilesSource{
		packageID: pkg,
		files:     files,
		curIdx:    -1,
		length:    len(files),
		err:       nil,
	}

	return tx.CopyFrom(ctx, pgx.Identifier{"mod_package_files"}, filesSourceCols, &source)
}

func main() {
	cfg, loader := config.Loader()
	// flags := loader.Flags()

	if err := loader.Load(); err != nil {
		if strings.Contains(err.Error(), "help requested") {
			os.Exit(3)
		}

		panic(err)
	}

	writer := zerolog.ConsoleWriter{Out: os.Stdout}
	writer.TimeFormat = "02.01.2006 15:04:05 MST"
	writer.FormatFieldValue = func(value interface{}) string {
		if value == nil {
			return "                "
		}

		str, ok := value.(string)
		if ok {
			if len(str) == 16 {
				// color request IDs in cyan  we have to guess based on the field content because we can't get
				// the current field name
				return fmt.Sprintf("\x1b[%dm%s\x1b[0m", 36, value)
			} else if strings.Contains(str, "\\n") && strings.Contains(str, "\\t") {
				// unquote values that contain line breaks and tabs because they're most likely stack traces
				str, err := strconv.Unquote(str)
				if err == nil {
					return str
				}
			}
		}

		return fmt.Sprintf("%s", value)
	}

	log.Logger = log.Output(writer)

	zerolog.ErrorStackMarshaler = func(err error) interface{} {
		return eris.ToString(err, true)
	}

	if err := cfg.Validate(); err != nil {
		log.Fatal().Err(err).Msg("Failed to parse config")
	}

	zerolog.SetGlobalLevel(cfg.LogLevel())
	log.Logger = log.Logger.With().Caller().Stack().Logger()

	dbConfig, err := pgxpool.ParseConfig(cfg.Database)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to parse DB config")
	}

	dbConfig.ConnConfig.Logger = nblog.PgxLogger{}
	dbConfig.ConnConfig.LogLevel = pgx.LogLevelDebug
	pool, err := pgxpool.ConnectConfig(context.Background(), dbConfig)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to open DB connection")
	}

	log.Info().Msg("Finished parsing configuration; loading data")

	inFile, err := os.Open("repo.json")
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to open repo.json")
	}
	defer inFile.Close()

	jsonStr, err := ioutil.ReadAll(inFile)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to read repo.json")
	}

	data := importer.KnStruct{}
	err = json.Unmarshal(jsonStr, &data)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to parse repo.json")
	}

	typeMap := map[string]db.ModType{
		"mod":    db.TypeMod,
		"tc":     db.TypeTotalConversion,
		"engine": db.TypeEngine,
		"tool":   db.TypeTool,
		"ext":    db.TypeExtension,
	}

	stabilityMap := map[string]db.EngineStability{
		"":        db.EngineUnknown,
		"stable":  db.EngineStable,
		"rc":      db.EngineRC,
		"nightly": db.EngineNightly,
	}

	pkgTypeMap := map[string]db.PackageType{
		"required":    db.PackageRequired,
		"recommended": db.PackageRecommended,
		"optional":    db.PackageOptional,
	}

	ctx := context.Background()
	tx, err := pool.Begin(ctx)
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to initiate transaction")
	}

	_, err = tx.Exec(ctx, "SET CONSTRAINTS ALL DEFERRED")
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to setup the transaction")
	}

	q := queries.NewQuerier(tx)

	count := len(data.Mods)
	aidCache := map[string]int32{}
	for i, mod := range data.Mods {
		log.Info().Msgf("Mod %4d of %d", i, count)

		aid, ok := aidCache[mod.ID]
		if !ok {
			mType, ok := typeMap[mod.Type]
			if !ok {
				log.Warn().Msgf("Invalid mod type %s detected for mod %s (%s)", mod.Type, mod.Title, mod.ID)
				mType = db.TypeMod
			}

			aid, err = q.CreateMod(ctx, queries.CreateModParams{
				Modid: mod.ID,
				Title: mod.Title,
				Type:  int16(mType),
			})
			if err != nil {
				log.Fatal().Err(err).Msgf("Failed to insert mod %s (%s)", mod.Title, mod.ID)
			}

			aidCache[mod.ID] = aid
		}

		log.Info().Msgf("AID %d", aid)

		rStabilility, ok := stabilityMap[mod.Stability]
		if !ok {
			log.Warn().Msgf("Invalid stability %s detected for mod %s (%s)", mod.Stability, mod.ID, mod.Version)
			rStabilility = db.EngineUnknown
		}

		// TODO files (logo -> ?, tile -> teaser, banner -> banner, screenshots)
		relDate, err := time.Parse("2006-01-02", mod.FirstRelease)
		if err != nil {
			if mod.FirstRelease == "" {
				relDate = time.Unix(1, 0)
			} else {
				log.Fatal().Msgf("Invalid release date for mod %s (%s): %s", mod.ID, mod.Version, mod.FirstRelease)
			}
		}

		updateDate, err := time.Parse("2006-01-02", mod.LastUpdate)
		if err != nil {
			if mod.LastUpdate == "" {
				updateDate = time.Unix(1, 0)
			} else {
				log.Fatal().Msgf("Invalid update date for mod %s (%s): %s", mod.ID, mod.Version, mod.LastUpdate)
			}
		}

		screens := make([]int32, len(mod.Screenshots))
		for idx, url := range mod.Screenshots {
			screens[idx] = handleFile(ctx, q, url)
		}

		modParams := queries.CreateReleaseParams{
			ModAid:        aid,
			Version:       mod.Version,
			Stability:     int16(rStabilility),
			Description:   mod.Description,
			ReleaseThread: mod.ReleaseThread,
			Videos:        mod.Videos,
			Notes:         mod.Notes,
			Cmdline:       mod.Cmdline,
			Screenshots:   screens,
			Teaser:        handleFile(ctx, q, mod.Tile),
			Banner:        handleFile(ctx, q, mod.Banner),
		}

		modParams.Released.Set(relDate)
		modParams.Updated.Set(updateDate)

		rid, err := q.CreateRelease(ctx, modParams)
		if err != nil {
			log.Fatal().Err(err).Msgf("Failed to insert release %s (%s)", mod.ID, mod.Version)
		}

		for _, pkg := range mod.Packages {
			pkgType, ok := pkgTypeMap[pkg.Status]
			if !ok {
				log.Fatal().Msgf("Failed to parse package status %s for %s in %s (%s)", pkg.Status, pkg.Name, mod.ID, mod.Version)
			}

			pid, err := q.CreatePackage(ctx, queries.CreatePackageParams{
				ReleaseID: rid,
				Name:      pkg.Name,
				Folder:    pkg.Folder,
				Notes:     pkg.Notes,
				Type:      int16(pkgType),
				// TODO: Actually parse the env spec
				CpuSpecs:  strings.Split(pkg.Environment, " && "),
				KnossosVp: pkg.IsVp,
			})
			if err != nil {
				log.Fatal().Err(err).Msgf("Failed to insert package %s for %s (%s)", pkg.Name, mod.ID, mod.Version)
			}

			pkgBatch := new(pgx.Batch)

			for _, dep := range pkg.Dependencies {
				q.CreatePackageDependencyBatch(pkgBatch, queries.CreatePackageDependencyParams{
					PackageID: pid,
					Modid:     dep.ID,
					Version:   dep.Version,
					Packages:  dep.Packages,
				})
			}

			for _, exe := range pkg.Executables {
				prio := int16(0)
				// See https://github.com/ngld/knossos/blob/1f60d925498c02d3db76a54d3ee20c31b75c5a21/knossos/repo.py#L35-L40
				if exe.Properties.X64 {
					prio += 50
				}
				if exe.Properties.AVX2 {
					prio += 3
				}
				if exe.Properties.AVX {
					prio += 2
				}
				if exe.Properties.SSE2 {
					prio++
				}

				q.CreatePackageExecutableBatch(pkgBatch, queries.CreatePackageExecutableParams{
					PackageID: pid,
					Path:      exe.File,
					Label:     exe.Label,
					Priority:  prio,
					Debug:     strings.Contains(exe.Label, "Debug"),
				})
			}

			for _, archive := range pkg.Files {
				digest, err := hex.DecodeString(archive.Checksum[1])
				if err != nil {
					log.Fatal().Err(err).Msgf("Failed to parse digest for archive %s of %s for %s (%s)", archive.Filename, pkg.Name, mod.ID, mod.Version)
				}

				fid, err := q.CreateExternalFile(ctx, queries.CreateExternalFileParams{
					StorageKey: "ext#" + nanoid.New(),
					Filesize:   int32(archive.FileSize),
					Public:     true,
					External:   archive.URLs,
					Owner:      1,
				})
				if err != nil {
					log.Fatal().Err(err).Msgf("Failed to index file for archive %s of %s for %s (%s)", archive.Filename, pkg.Name, mod.ID, mod.Version)
				}

				q.CreatePackageArchiveBatch(pkgBatch, queries.CreatePackageArchiveParams{
					PackageID:      pid,
					Label:          archive.Filename,
					Destination:    archive.Dest,
					ChecksumAlgo:   archive.Checksum[0],
					ChecksumDigest: digest,
					FileID:         fid,
				})
			}

			count, err := CopyFromFiles(ctx, tx, pid, pkg.Filelist)
			if err != nil {
				log.Fatal().Err(err).Msgf("Failed to process filelist on %s for %s (%s)", pkg.Name, mod.ID, mod.Version)
			}
			if count != int64(len(pkg.Filelist)) {
				log.Fatal().Err(err).Msgf("Imported only %d of %d files on %s for %s (%s)", count, len(pkg.Filelist), pkg.Name, mod.ID, mod.Version)
			}

			batchResults := tx.SendBatch(ctx, pkgBatch)
			for idx := 0; idx < pkgBatch.Len(); idx++ {
				_, err := batchResults.Exec()
				if err != nil {
					log.Fatal().Err(err).Msgf("Failed process package contents for %s of %s (%s)", pkg.Name, mod.ID, mod.Version)
				}
			}
			batchResults.Close()
		}
	}

	log.Info().Msg("Comitting")
	err = tx.Commit(ctx)
	if err != nil {
		log.Fatal().Err(err).Msg("Commit failed")
	}

	log.Info().Msg("Done")
}
<|endoftext|>


<commit_before>package xmlquery

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func findNode(root *Node, name string) *Node {
	node := root.FirstChild
	for {
		if node == nil || node.Data == name {
			break
		}
		node = node.NextSibling
	}
	return node
}

func childNodes(root *Node, name string) []*Node {
	var list []*Node
	node := root.FirstChild
	for {
		if node == nil {
			break
		}
		if node.Data == name {
			list = append(list, node)
		}
		node = node.NextSibling
	}
	return list
}

func testNode(t *testing.T, n *Node, expected string) {
	if n.Data != expected {
		t.Fatalf("expected node name is %s,but got %s", expected, n.Data)
	}
}

func testAttr(t *testing.T, n *Node, name, expected string) {
	for _, attr := range n.Attr {
		if attr.Name.Local == name && attr.Value == expected {
			return
		}
	}
	t.Fatalf("not found attribute %s in the node %s", name, n.Data)
}

func testValue(t *testing.T, val, expected string) {
	if val != expected {
		t.Fatalf("expected value is %s,but got %s", expected, val)
	}
}

func TestLoadURL(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s := `<?xml version="1.0"?>
       <rss>
	   	<title></title>
	   </rss>`
		w.Header().Set("Content-Type", "text/xml")
		w.Write([]byte(s))
	}))
	defer server.Close()
	_, err := LoadURL(server.URL)
	if err != nil {
		t.Fatal(err)
	}
}

func TestNamespaceURL(t *testing.T) {
	s := `
<?xml version="1.0"?>
<rss version="2.0" xmlns:dc="https://purl.org/dc/elements/1.1/">
<!-- author -->
<dc:creator><![CDATA[Richard Lawler]]></dc:creator>
<dc:identifier>21|22021348</dc:identifier>
</rss>
	`
	doc, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Fatal(err)
	}
	top := FindOne(doc, "//rss")
	if top == nil {
		t.Fatal("rss feed invalid")
	}
	node := FindOne(top, "dc:creator")
	if node.Prefix != "dc" {
		t.Fatalf("expected node prefix name is dc but is=%s", node.Prefix)
	}
	if node.NamespaceURI != "https://purl.org/dc/elements/1.1/" {
		t.Fatalf("dc:creator != %s", node.NamespaceURI)
	}
	if strings.Index(top.InnerText(), "author") > 0 {
		t.Fatalf("InnerText() include comment node text")
	}
	if strings.Index(top.OutputXML(true), "author") == -1 {
		t.Fatal("OutputXML shoud include comment node,but not")
	}
}

func TestMultipleProcInst(t *testing.T) {
	s := `<?xml version="1.0" encoding="UTF-8"?>
<?xml-stylesheet type="text/xsl" media="screen" href="/~d/styles/rss2full.xsl"?>
<?xml-stylesheet type="text/css" media="screen" href="http://feeds.reuters.com/~d/styles/itemcontent.css"?>
<rss xmlns:feedburner="http://rssnamespace.org/feedburner/ext/1.0" version="2.0">
</rss>
	`
	doc, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Fatal(err)
	}

	node := doc.FirstChild // <?xml ?>
	if node.Data != "xml" {
		t.Fatal("node.Data != xml")
	}
	node = node.NextSibling // New Line
	node = node.NextSibling // <?xml-stylesheet?>
	if node.Data != "xml-stylesheet" {
		t.Fatal("node.Data != xml-stylesheet")
	}
}

func TestParse(t *testing.T) {
	s := `<?xml version="1.0" encoding="UTF-8"?>
<bookstore>
<book>
  <title lang="en">Harry Potter</title>
  <price>29.99</price>
</book>
<book>
  <title lang="en">Learning XML</title>
  <price>39.95</price>
</book>
</bookstore>`
	root, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Error(err)
	}
	if root.Type != DocumentNode {
		t.Fatal("top node of tree is not DocumentNode")
	}

	declarNode := root.FirstChild
	if declarNode.Type != DeclarationNode {
		t.Fatal("first child node of tree is not DeclarationNode")
	}

	if declarNode.Attr[0].Name.Local != "version" && declarNode.Attr[0].Value != "1.0" {
		t.Fatal("version attribute not expected")
	}

	bookstore := root.LastChild
	if bookstore.Data != "bookstore" {
		t.Fatal("bookstore elem not found")
	}
	if bookstore.FirstChild.Data != "\n" {
		t.Fatal("first child node of bookstore is not empty node(\n)")
	}
	books := childNodes(bookstore, "book")
	if len(books) != 2 {
		t.Fatalf("expected book element count is 2, but got %d", len(books))
	}
	// first book element
	testNode(t, findNode(books[0], "title"), "title")
	testAttr(t, findNode(books[0], "title"), "lang", "en")
	testValue(t, findNode(books[0], "price").InnerText(), "29.99")
	testValue(t, findNode(books[0], "title").InnerText(), "Harry Potter")

	// second book element
	testNode(t, findNode(books[1], "title"), "title")
	testAttr(t, findNode(books[1], "title"), "lang", "en")
	testValue(t, findNode(books[1], "price").InnerText(), "39.95")

	testValue(t, books[0].OutputXML(true), `<book><title lang="en">Harry Potter</title><price>29.99</price></book>`)
}

func TestMissDeclaration(t *testing.T) {
	s := `<AAA>
		<BBB></BBB>
		<CCC></CCC>
	</AAA>`
	doc, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Fatal(err)
	}
	node := FindOne(doc, "//AAA")
	if node == nil {
		t.Fatal("//AAA is nil")
	}
}

func TestTooNested(t *testing.T) {
	s := `<?xml version="1.0" encoding="UTF-8"?>
	<!-- comment here-->
    <AAA>
        <BBB>
            <DDD>
                <CCC>
                    <DDD/>
                    <EEE/>
                </CCC>
            </DDD>
        </BBB>
        <CCC>
            <DDD>
                <EEE>
                    <DDD>
                        <FFF/>
                    </DDD>
                </EEE>
            </DDD>
        </CCC>
     </AAA>`
	root, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Error(err)
	}
	aaa := findNode(root, "AAA")
	if aaa == nil {
		t.Fatal("AAA node not exists")
	}
	ccc := aaa.LastChild
	if ccc.Data != "CCC" {
		t.Fatalf("expected node is CCC,but got %s", ccc.Data)
	}
	bbb := ccc.PrevSibling
	if bbb.Data != "BBB" {
		t.Fatalf("expected node is bbb,but got %s", bbb.Data)
	}
	ddd := findNode(bbb, "DDD")
	testNode(t, ddd, "DDD")
	testNode(t, ddd.LastChild, "CCC")
}

func TestSelectElement(t *testing.T) {
	s := `<?xml version="1.0" encoding="UTF-8"?>
    <AAA>
        <BBB id="1"/>
        <CCC id="2">
            <DDD/>
        </CCC>
		<CCC id="3">
            <DDD/>
        </CCC>
     </AAA>`
	root, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Error(err)
	}
	version := root.FirstChild.SelectAttr("version")
	if version != "1.0" {
		t.Fatal("version!=1.0")
	}
	aaa := findNode(root, "AAA")
	var n *Node
	n = aaa.SelectElement("BBB")
	if n == nil {
		t.Fatalf("n is nil")
	}
	n = aaa.SelectElement("CCC")
	if n == nil {
		t.Fatalf("n is nil")
	}

	var ns []*Node
	ns = aaa.SelectElements("CCC")
	if len(ns) != 2 {
		t.Fatalf("len(ns)!=2")
	}
}

func TestEscapeOutputValue(t *testing.T) {
	data := `<AAA><*></AAA>`
}
<commit_msg>Add test<commit_after>package xmlquery

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func findNode(root *Node, name string) *Node {
	node := root.FirstChild
	for {
		if node == nil || node.Data == name {
			break
		}
		node = node.NextSibling
	}
	return node
}

func childNodes(root *Node, name string) []*Node {
	var list []*Node
	node := root.FirstChild
	for {
		if node == nil {
			break
		}
		if node.Data == name {
			list = append(list, node)
		}
		node = node.NextSibling
	}
	return list
}

func testNode(t *testing.T, n *Node, expected string) {
	if n.Data != expected {
		t.Fatalf("expected node name is %s,but got %s", expected, n.Data)
	}
}

func testAttr(t *testing.T, n *Node, name, expected string) {
	for _, attr := range n.Attr {
		if attr.Name.Local == name && attr.Value == expected {
			return
		}
	}
	t.Fatalf("not found attribute %s in the node %s", name, n.Data)
}

func testValue(t *testing.T, val, expected string) {
	if val != expected {
		t.Fatalf("expected value is %s,but got %s", expected, val)
	}
}

func TestLoadURL(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s := `<?xml version="1.0"?>
       <rss>
	   	<title></title>
	   </rss>`
		w.Header().Set("Content-Type", "text/xml")
		w.Write([]byte(s))
	}))
	defer server.Close()
	_, err := LoadURL(server.URL)
	if err != nil {
		t.Fatal(err)
	}
}

func TestNamespaceURL(t *testing.T) {
	s := `
<?xml version="1.0"?>
<rss version="2.0" xmlns:dc="https://purl.org/dc/elements/1.1/">
<!-- author -->
<dc:creator><![CDATA[Richard Lawler]]></dc:creator>
<dc:identifier>21|22021348</dc:identifier>
</rss>
	`
	doc, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Fatal(err)
	}
	top := FindOne(doc, "//rss")
	if top == nil {
		t.Fatal("rss feed invalid")
	}
	node := FindOne(top, "dc:creator")
	if node.Prefix != "dc" {
		t.Fatalf("expected node prefix name is dc but is=%s", node.Prefix)
	}
	if node.NamespaceURI != "https://purl.org/dc/elements/1.1/" {
		t.Fatalf("dc:creator != %s", node.NamespaceURI)
	}
	if strings.Index(top.InnerText(), "author") > 0 {
		t.Fatalf("InnerText() include comment node text")
	}
	if strings.Index(top.OutputXML(true), "author") == -1 {
		t.Fatal("OutputXML shoud include comment node,but not")
	}
}

func TestMultipleProcInst(t *testing.T) {
	s := `<?xml version="1.0" encoding="UTF-8"?>
<?xml-stylesheet type="text/xsl" media="screen" href="/~d/styles/rss2full.xsl"?>
<?xml-stylesheet type="text/css" media="screen" href="http://feeds.reuters.com/~d/styles/itemcontent.css"?>
<rss xmlns:feedburner="http://rssnamespace.org/feedburner/ext/1.0" version="2.0">
</rss>
	`
	doc, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Fatal(err)
	}

	node := doc.FirstChild // <?xml ?>
	if node.Data != "xml" {
		t.Fatal("node.Data != xml")
	}
	node = node.NextSibling // New Line
	node = node.NextSibling // <?xml-stylesheet?>
	if node.Data != "xml-stylesheet" {
		t.Fatal("node.Data != xml-stylesheet")
	}
}

func TestParse(t *testing.T) {
	s := `<?xml version="1.0" encoding="UTF-8"?>
<bookstore>
<book>
  <title lang="en">Harry Potter</title>
  <price>29.99</price>
</book>
<book>
  <title lang="en">Learning XML</title>
  <price>39.95</price>
</book>
</bookstore>`
	root, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Error(err)
	}
	if root.Type != DocumentNode {
		t.Fatal("top node of tree is not DocumentNode")
	}

	declarNode := root.FirstChild
	if declarNode.Type != DeclarationNode {
		t.Fatal("first child node of tree is not DeclarationNode")
	}

	if declarNode.Attr[0].Name.Local != "version" && declarNode.Attr[0].Value != "1.0" {
		t.Fatal("version attribute not expected")
	}

	bookstore := root.LastChild
	if bookstore.Data != "bookstore" {
		t.Fatal("bookstore elem not found")
	}
	if bookstore.FirstChild.Data != "\n" {
		t.Fatal("first child node of bookstore is not empty node(\n)")
	}
	books := childNodes(bookstore, "book")
	if len(books) != 2 {
		t.Fatalf("expected book element count is 2, but got %d", len(books))
	}
	// first book element
	testNode(t, findNode(books[0], "title"), "title")
	testAttr(t, findNode(books[0], "title"), "lang", "en")
	testValue(t, findNode(books[0], "price").InnerText(), "29.99")
	testValue(t, findNode(books[0], "title").InnerText(), "Harry Potter")

	// second book element
	testNode(t, findNode(books[1], "title"), "title")
	testAttr(t, findNode(books[1], "title"), "lang", "en")
	testValue(t, findNode(books[1], "price").InnerText(), "39.95")

	testValue(t, books[0].OutputXML(true), `<book><title lang="en">Harry Potter</title><price>29.99</price></book>`)
}

func TestMissDeclaration(t *testing.T) {
	s := `<AAA>
		<BBB></BBB>
		<CCC></CCC>
	</AAA>`
	doc, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Fatal(err)
	}
	node := FindOne(doc, "//AAA")
	if node == nil {
		t.Fatal("//AAA is nil")
	}
}

func TestTooNested(t *testing.T) {
	s := `<?xml version="1.0" encoding="UTF-8"?>
	<!-- comment here-->
    <AAA>
        <BBB>
            <DDD>
                <CCC>
                    <DDD/>
                    <EEE/>
                </CCC>
            </DDD>
        </BBB>
        <CCC>
            <DDD>
                <EEE>
                    <DDD>
                        <FFF/>
                    </DDD>
                </EEE>
            </DDD>
        </CCC>
     </AAA>`
	root, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Error(err)
	}
	aaa := findNode(root, "AAA")
	if aaa == nil {
		t.Fatal("AAA node not exists")
	}
	ccc := aaa.LastChild
	if ccc.Data != "CCC" {
		t.Fatalf("expected node is CCC,but got %s", ccc.Data)
	}
	bbb := ccc.PrevSibling
	if bbb.Data != "BBB" {
		t.Fatalf("expected node is bbb,but got %s", bbb.Data)
	}
	ddd := findNode(bbb, "DDD")
	testNode(t, ddd, "DDD")
	testNode(t, ddd.LastChild, "CCC")
}

func TestSelectElement(t *testing.T) {
	s := `<?xml version="1.0" encoding="UTF-8"?>
    <AAA>
        <BBB id="1"/>
        <CCC id="2">
            <DDD/>
        </CCC>
		<CCC id="3">
            <DDD/>
        </CCC>
     </AAA>`
	root, err := Parse(strings.NewReader(s))
	if err != nil {
		t.Error(err)
	}
	version := root.FirstChild.SelectAttr("version")
	if version != "1.0" {
		t.Fatal("version!=1.0")
	}
	aaa := findNode(root, "AAA")
	var n *Node
	n = aaa.SelectElement("BBB")
	if n == nil {
		t.Fatalf("n is nil")
	}
	n = aaa.SelectElement("CCC")
	if n == nil {
		t.Fatalf("n is nil")
	}

	var ns []*Node
	ns = aaa.SelectElements("CCC")
	if len(ns) != 2 {
		t.Fatalf("len(ns)!=2")
	}
}

func TestEscapeOutputValue(t *testing.T) {
	data := `<AAA><*></AAA>`

	root, err := Parse(strings.NewReader(data))
	if err != nil {
		t.Error(err)
	}

	escapedInnerText := root.OutputXML(true)
	if !strings.Contains(escapedInnerText, "<*>") {
		t.Fatal("Inner Text has not been escaped")
	}

}
<|endoftext|>


<commit_before>// Data structs (and functions to act on them) to hold Player information

package main

import "fmt"

type Gender int

const (
	Male Gender = iota
	Female
	Default
)

// StringToGender parses single character string into a Gender.
//
// The string must be either "m" or "f", or we return error.
func StringToGender(s string) (Gender, error) {
	switch s {
	case "m":
		return Male, nil
	case "f":
		return Female, nil
	}
	return Default, fmt.Errorf("invalid gender '%s'", s)
}

type Player struct {
	name   string
	value  float32
	gender Gender
	team   uint8
}
<commit_msg>Gender type fix<commit_after>// Data structs (and functions to act on them) to hold Player information

package main

import "fmt"

type Gender uint8

const (
	Male Gender = iota
	Female
	Default
)

// StringToGender parses single character string into a Gender.
//
// The string must be either "m" or "f", or we return error.
func StringToGender(s string) (Gender, error) {
	switch s {
	case "m":
		return Male, nil
	case "f":
		return Female, nil
	}
	return Default, fmt.Errorf("invalid gender '%s'", s)
}

type Player struct {
	name   string
	value  float32
	gender Gender
	team   uint8
}
<|endoftext|>


<commit_before>//
// Copyright 2014-2020 Cristian Maglie. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package enumerator

import (
	"fmt"
	"regexp"
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows"
)

func parseDeviceID(deviceID string, details *PortDetails) {
	// Windows stock USB-CDC driver
	if len(deviceID) >= 3 && deviceID[:3] == "USB" {
		re := regexp.MustCompile("VID_(....)&PID_(....)(\\\\(\\w+)$)?").FindAllStringSubmatch(deviceID, -1)
		if re == nil || len(re[0]) < 2 {
			// Silently ignore unparsable strings
			return
		}
		details.IsUSB = true
		details.VID = re[0][1]
		details.PID = re[0][2]
		if len(re[0]) >= 4 {
			details.SerialNumber = re[0][4]
		}
		return
	}

	// FTDI driver
	if len(deviceID) >= 7 && deviceID[:7] == "FTDIBUS" {
		re := regexp.MustCompile("VID_(....)\\+PID_(....)(\\+(\\w+))?").FindAllStringSubmatch(deviceID, -1)
		if re == nil || len(re[0]) < 2 {
			// Silently ignore unparsable strings
			return
		}
		details.IsUSB = true
		details.VID = re[0][1]
		details.PID = re[0][2]
		if len(re[0]) >= 4 {
			details.SerialNumber = re[0][4]
		}
		return
	}

	// Other unidentified device type
}

// setupapi based
// --------------

//sys setupDiClassGuidsFromNameInternal(class string, guid *guid, guidSize uint32, requiredSize *uint32) (err error) = setupapi.SetupDiClassGuidsFromNameW
//sys setupDiGetClassDevs(guid *guid, enumerator *string, hwndParent uintptr, flags uint32) (set devicesSet, err error) = setupapi.SetupDiGetClassDevsW
//sys setupDiDestroyDeviceInfoList(set devicesSet) (err error) = setupapi.SetupDiDestroyDeviceInfoList
//sys setupDiEnumDeviceInfo(set devicesSet, index uint32, info *devInfoData) (err error) = setupapi.SetupDiEnumDeviceInfo
//sys setupDiGetDeviceInstanceId(set devicesSet, devInfo *devInfoData, devInstanceId unsafe.Pointer, devInstanceIdSize uint32, requiredSize *uint32) (err error) = setupapi.SetupDiGetDeviceInstanceIdW
//sys setupDiOpenDevRegKey(set devicesSet, devInfo *devInfoData, scope dicsScope, hwProfile uint32, keyType uint32, samDesired regsam) (hkey syscall.Handle, err error) = setupapi.SetupDiOpenDevRegKey
//sys setupDiGetDeviceRegistryProperty(set devicesSet, devInfo *devInfoData, property deviceProperty, propertyType *uint32, outValue *byte, outSize *uint32, reqSize *uint32) (res bool) = setupapi.SetupDiGetDeviceRegistryPropertyW

// Device registry property codes
// (Codes marked as read-only (R) may only be used for
// SetupDiGetDeviceRegistryProperty)
//
// These values should cover the same set of registry properties
// as defined by the CM_DRP codes in cfgmgr32.h.
//
// Note that SPDRP codes are zero based while CM_DRP codes are one based!
type deviceProperty uint32

const (
	spdrpDeviceDesc               deviceProperty = 0x00000000 // DeviceDesc = R/W
	spdrpHardwareID                              = 0x00000001 // HardwareID = R/W
	spdrpCompatibleIDS                           = 0x00000002 // CompatibleIDs = R/W
	spdrpUnused0                                 = 0x00000003 // Unused
	spdrpService                                 = 0x00000004 // Service = R/W
	spdrpUnused1                                 = 0x00000005 // Unused
	spdrpUnused2                                 = 0x00000006 // Unused
	spdrpClass                                   = 0x00000007 // Class = R--tied to ClassGUID
	spdrpClassGUID                               = 0x00000008 // ClassGUID = R/W
	spdrpDriver                                  = 0x00000009 // Driver = R/W
	spdrpConfigFlags                             = 0x0000000A // ConfigFlags = R/W
	spdrpMFG                                     = 0x0000000B // Mfg = R/W
	spdrpFriendlyName                            = 0x0000000C // FriendlyName = R/W
	spdrpLocationIinformation                    = 0x0000000D // LocationInformation = R/W
	spdrpPhysicalDeviceObjectName                = 0x0000000E // PhysicalDeviceObjectName = R
	spdrpCapabilities                            = 0x0000000F // Capabilities = R
	spdrpUINumber                                = 0x00000010 // UiNumber = R
	spdrpUpperFilters                            = 0x00000011 // UpperFilters = R/W
	spdrpLowerFilters                            = 0x00000012 // LowerFilters = R/W
	spdrpBusTypeGUID                             = 0x00000013 // BusTypeGUID = R
	spdrpLegactBusType                           = 0x00000014 // LegacyBusType = R
	spdrpBusNumber                               = 0x00000015 // BusNumber = R
	spdrpEnumeratorName                          = 0x00000016 // Enumerator Name = R
	spdrpSecurity                                = 0x00000017 // Security = R/W, binary form
	spdrpSecuritySDS                             = 0x00000018 // Security = W, SDS form
	spdrpDevType                                 = 0x00000019 // Device Type = R/W
	spdrpExclusive                               = 0x0000001A // Device is exclusive-access = R/W
	spdrpCharacteristics                         = 0x0000001B // Device Characteristics = R/W
	spdrpAddress                                 = 0x0000001C // Device Address = R
	spdrpUINumberDescFormat                      = 0X0000001D // UiNumberDescFormat = R/W
	spdrpDevicePowerData                         = 0x0000001E // Device Power Data = R
	spdrpRemovalPolicy                           = 0x0000001F // Removal Policy = R
	spdrpRemovalPolicyHWDefault                  = 0x00000020 // Hardware Removal Policy = R
	spdrpRemovalPolicyOverride                   = 0x00000021 // Removal Policy Override = RW
	spdrpInstallState                            = 0x00000022 // Device Install State = R
	spdrpLocationPaths                           = 0x00000023 // Device Location Paths = R
	spdrpBaseContainerID                         = 0x00000024 // Base ContainerID = R

	spdrpMaximumProperty = 0x00000025 // Upper bound on ordinals
)

// Values specifying the scope of a device property change
type dicsScope uint32

const (
	dicsFlagGlobal          dicsScope = 0x00000001 // make change in all hardware profiles
	dicsFlagConfigSspecific           = 0x00000002 // make change in specified profile only
	dicsFlagConfigGeneral             = 0x00000004 // 1 or more hardware profile-specific
)

// https://msdn.microsoft.com/en-us/library/windows/desktop/ms724878(v=vs.85).aspx
type regsam uint32

const (
	keyAllAccess        regsam = 0xF003F
	keyCreateLink              = 0x00020
	keyCreateSubKey            = 0x00004
	keyEnumerateSubKeys        = 0x00008
	keyExecute                 = 0x20019
	keyNotify                  = 0x00010
	keyQueryValue              = 0x00001
	keyRead                    = 0x20019
	keySetValue                = 0x00002
	keyWOW64_32key             = 0x00200
	keyWOW64_64key             = 0x00100
	keyWrite                   = 0x20006
)

// KeyType values for SetupDiCreateDevRegKey, SetupDiOpenDevRegKey, and
// SetupDiDeleteDevRegKey.
const (
	diregDev  = 0x00000001 // Open/Create/Delete device key
	diregDrv  = 0x00000002 // Open/Create/Delete driver key
	diregBoth = 0x00000004 // Delete both driver and Device key
)

// https://msdn.microsoft.com/it-it/library/windows/desktop/aa373931(v=vs.85).aspx
type guid struct {
	data1 uint32
	data2 uint16
	data3 uint16
	data4 [8]byte
}

func (g guid) String() string {
	return fmt.Sprintf("%08x-%04x-%04x-%02x%02x-%02x%02x%02x%02x%02x%02x",
		g.data1, g.data2, g.data3,
		g.data4[0], g.data4[1], g.data4[2], g.data4[3],
		g.data4[4], g.data4[5], g.data4[6], g.data4[7])
}

func classGuidsFromName(className string) ([]guid, error) {
	// Determine the number of GUIDs for className
	n := uint32(0)
	if err := setupDiClassGuidsFromNameInternal(className, nil, 0, &n); err != nil {
		// ignore error: UIDs array size too small
	}

	res := make([]guid, n)
	err := setupDiClassGuidsFromNameInternal(className, &res[0], n, &n)
	return res, err
}

const (
	digcfDefault         = 0x00000001 // only valid with digcfDeviceInterface
	digcfPresent         = 0x00000002
	digcfAllClasses      = 0x00000004
	digcfProfile         = 0x00000008
	digcfDeviceInterface = 0x00000010
)

type devicesSet syscall.Handle

func (g *guid) getDevicesSet() (devicesSet, error) {
	return setupDiGetClassDevs(g, nil, 0, digcfPresent)
}

func (set devicesSet) destroy() {
	setupDiDestroyDeviceInfoList(set)
}

// https://msdn.microsoft.com/en-us/library/windows/hardware/ff552344(v=vs.85).aspx
type devInfoData struct {
	size     uint32
	guid     guid
	devInst  uint32
	reserved uintptr
}

type deviceInfo struct {
	set  devicesSet
	data devInfoData
}

func (set devicesSet) getDeviceInfo(index int) (*deviceInfo, error) {
	result := &deviceInfo{set: set}

	result.data.size = uint32(unsafe.Sizeof(result.data))
	err := setupDiEnumDeviceInfo(set, uint32(index), &result.data)
	return result, err
}

func (dev *deviceInfo) getInstanceID() (string, error) {
	n := uint32(0)
	setupDiGetDeviceInstanceId(dev.set, &dev.data, nil, 0, &n)
	buff := make([]uint16, n)
	if err := setupDiGetDeviceInstanceId(dev.set, &dev.data, unsafe.Pointer(&buff[0]), uint32(len(buff)), &n); err != nil {
		return "", err
	}
	return windows.UTF16ToString(buff[:]), nil
}

func (dev *deviceInfo) openDevRegKey(scope dicsScope, hwProfile uint32, keyType uint32, samDesired regsam) (syscall.Handle, error) {
	return setupDiOpenDevRegKey(dev.set, &dev.data, scope, hwProfile, keyType, samDesired)
}

func nativeGetDetailedPortsList() ([]*PortDetails, error) {
	guids, err := classGuidsFromName("Ports")
	if err != nil {
		return nil, &PortEnumerationError{causedBy: err}
	}

	var res []*PortDetails
	for _, g := range guids {
		devsSet, err := g.getDevicesSet()
		if err != nil {
			return nil, &PortEnumerationError{causedBy: err}
		}
		defer devsSet.destroy()

		for i := 0; ; i++ {
			device, err := devsSet.getDeviceInfo(i)
			if err != nil {
				break
			}
			details := &PortDetails{}
			portName, err := retrievePortNameFromDevInfo(device)
			if err != nil {
				continue
			}
			if len(portName) < 3 || portName[0:3] != "COM" {
				// Accept only COM ports
				continue
			}
			details.Name = portName

			if err := retrievePortDetailsFromDevInfo(device, details); err != nil {
				return nil, &PortEnumerationError{causedBy: err}
			}
			res = append(res, details)
		}
	}
	return res, nil
}

func retrievePortNameFromDevInfo(device *deviceInfo) (string, error) {
	h, err := device.openDevRegKey(dicsFlagGlobal, 0, diregDev, keyRead)
	if err != nil {
		return "", err
	}
	defer syscall.RegCloseKey(h)

	var name [1024]uint16
	nameP := (*byte)(unsafe.Pointer(&name[0]))
	nameSize := uint32(len(name) * 2)
	if err := syscall.RegQueryValueEx(h, syscall.StringToUTF16Ptr("PortName"), nil, nil, nameP, &nameSize); err != nil {
		return "", err
	}
	return syscall.UTF16ToString(name[:]), nil
}

func retrievePortDetailsFromDevInfo(device *deviceInfo, details *PortDetails) error {
	deviceID, err := device.getInstanceID()
	if err != nil {
		return err
	}
	parseDeviceID(deviceID, details)

	var friendlyName [1024]uint16
	friendlyNameP := (*byte)(unsafe.Pointer(&friendlyName[0]))
	friendlyNameSize := uint32(len(friendlyName) * 2)
	if setupDiGetDeviceRegistryProperty(device.set, &device.data, spdrpDeviceDesc /* spdrpFriendlyName */, nil, friendlyNameP, &friendlyNameSize, nil) {
		//details.Product = syscall.UTF16ToString(friendlyName[:])
	}

	return nil
}
<commit_msg>completed retrievePortDetailsFromDevInfo<commit_after>//
// Copyright 2014-2020 Cristian Maglie. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package enumerator

import (
	"fmt"
	"regexp"
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows"
)

func parseDeviceID(deviceID string, details *PortDetails) {
	// Windows stock USB-CDC driver
	if len(deviceID) >= 3 && deviceID[:3] == "USB" {
		re := regexp.MustCompile("VID_(....)&PID_(....)(\\\\(\\w+)$)?").FindAllStringSubmatch(deviceID, -1)
		if re == nil || len(re[0]) < 2 {
			// Silently ignore unparsable strings
			return
		}
		details.IsUSB = true
		details.VID = re[0][1]
		details.PID = re[0][2]
		if len(re[0]) >= 4 {
			details.SerialNumber = re[0][4]
		}
		return
	}

	// FTDI driver
	if len(deviceID) >= 7 && deviceID[:7] == "FTDIBUS" {
		re := regexp.MustCompile("VID_(....)\\+PID_(....)(\\+(\\w+))?").FindAllStringSubmatch(deviceID, -1)
		if re == nil || len(re[0]) < 2 {
			// Silently ignore unparsable strings
			return
		}
		details.IsUSB = true
		details.VID = re[0][1]
		details.PID = re[0][2]
		if len(re[0]) >= 4 {
			details.SerialNumber = re[0][4]
		}
		return
	}

	// Other unidentified device type
}

// setupapi based
// --------------

//sys setupDiClassGuidsFromNameInternal(class string, guid *guid, guidSize uint32, requiredSize *uint32) (err error) = setupapi.SetupDiClassGuidsFromNameW
//sys setupDiGetClassDevs(guid *guid, enumerator *string, hwndParent uintptr, flags uint32) (set devicesSet, err error) = setupapi.SetupDiGetClassDevsW
//sys setupDiDestroyDeviceInfoList(set devicesSet) (err error) = setupapi.SetupDiDestroyDeviceInfoList
//sys setupDiEnumDeviceInfo(set devicesSet, index uint32, info *devInfoData) (err error) = setupapi.SetupDiEnumDeviceInfo
//sys setupDiGetDeviceInstanceId(set devicesSet, devInfo *devInfoData, devInstanceId unsafe.Pointer, devInstanceIdSize uint32, requiredSize *uint32) (err error) = setupapi.SetupDiGetDeviceInstanceIdW
//sys setupDiOpenDevRegKey(set devicesSet, devInfo *devInfoData, scope dicsScope, hwProfile uint32, keyType uint32, samDesired regsam) (hkey syscall.Handle, err error) = setupapi.SetupDiOpenDevRegKey
//sys setupDiGetDeviceRegistryProperty(set devicesSet, devInfo *devInfoData, property deviceProperty, propertyType *uint32, outValue *byte, bufSize uint32, reqSize *uint32) (res bool) = setupapi.SetupDiGetDeviceRegistryPropertyW

// Device registry property codes
// (Codes marked as read-only (R) may only be used for
// SetupDiGetDeviceRegistryProperty)
//
// These values should cover the same set of registry properties
// as defined by the CM_DRP codes in cfgmgr32.h.
//
// Note that SPDRP codes are zero based while CM_DRP codes are one based!
type deviceProperty uint32

const (
	spdrpDeviceDesc               deviceProperty = 0x00000000 // DeviceDesc = R/W
	spdrpHardwareID                              = 0x00000001 // HardwareID = R/W
	spdrpCompatibleIDS                           = 0x00000002 // CompatibleIDs = R/W
	spdrpUnused0                                 = 0x00000003 // Unused
	spdrpService                                 = 0x00000004 // Service = R/W
	spdrpUnused1                                 = 0x00000005 // Unused
	spdrpUnused2                                 = 0x00000006 // Unused
	spdrpClass                                   = 0x00000007 // Class = R--tied to ClassGUID
	spdrpClassGUID                               = 0x00000008 // ClassGUID = R/W
	spdrpDriver                                  = 0x00000009 // Driver = R/W
	spdrpConfigFlags                             = 0x0000000A // ConfigFlags = R/W
	spdrpMFG                                     = 0x0000000B // Mfg = R/W
	spdrpFriendlyName                            = 0x0000000C // FriendlyName = R/W
	spdrpLocationIinformation                    = 0x0000000D // LocationInformation = R/W
	spdrpPhysicalDeviceObjectName                = 0x0000000E // PhysicalDeviceObjectName = R
	spdrpCapabilities                            = 0x0000000F // Capabilities = R
	spdrpUINumber                                = 0x00000010 // UiNumber = R
	spdrpUpperFilters                            = 0x00000011 // UpperFilters = R/W
	spdrpLowerFilters                            = 0x00000012 // LowerFilters = R/W
	spdrpBusTypeGUID                             = 0x00000013 // BusTypeGUID = R
	spdrpLegactBusType                           = 0x00000014 // LegacyBusType = R
	spdrpBusNumber                               = 0x00000015 // BusNumber = R
	spdrpEnumeratorName                          = 0x00000016 // Enumerator Name = R
	spdrpSecurity                                = 0x00000017 // Security = R/W, binary form
	spdrpSecuritySDS                             = 0x00000018 // Security = W, SDS form
	spdrpDevType                                 = 0x00000019 // Device Type = R/W
	spdrpExclusive                               = 0x0000001A // Device is exclusive-access = R/W
	spdrpCharacteristics                         = 0x0000001B // Device Characteristics = R/W
	spdrpAddress                                 = 0x0000001C // Device Address = R
	spdrpUINumberDescFormat                      = 0X0000001D // UiNumberDescFormat = R/W
	spdrpDevicePowerData                         = 0x0000001E // Device Power Data = R
	spdrpRemovalPolicy                           = 0x0000001F // Removal Policy = R
	spdrpRemovalPolicyHWDefault                  = 0x00000020 // Hardware Removal Policy = R
	spdrpRemovalPolicyOverride                   = 0x00000021 // Removal Policy Override = RW
	spdrpInstallState                            = 0x00000022 // Device Install State = R
	spdrpLocationPaths                           = 0x00000023 // Device Location Paths = R
	spdrpBaseContainerID                         = 0x00000024 // Base ContainerID = R

	spdrpMaximumProperty = 0x00000025 // Upper bound on ordinals
)

// Values specifying the scope of a device property change
type dicsScope uint32

const (
	dicsFlagGlobal          dicsScope = 0x00000001 // make change in all hardware profiles
	dicsFlagConfigSspecific           = 0x00000002 // make change in specified profile only
	dicsFlagConfigGeneral             = 0x00000004 // 1 or more hardware profile-specific
)

// https://msdn.microsoft.com/en-us/library/windows/desktop/ms724878(v=vs.85).aspx
type regsam uint32

const (
	keyAllAccess        regsam = 0xF003F
	keyCreateLink              = 0x00020
	keyCreateSubKey            = 0x00004
	keyEnumerateSubKeys        = 0x00008
	keyExecute                 = 0x20019
	keyNotify                  = 0x00010
	keyQueryValue              = 0x00001
	keyRead                    = 0x20019
	keySetValue                = 0x00002
	keyWOW64_32key             = 0x00200
	keyWOW64_64key             = 0x00100
	keyWrite                   = 0x20006
)

// KeyType values for SetupDiCreateDevRegKey, SetupDiOpenDevRegKey, and
// SetupDiDeleteDevRegKey.
const (
	diregDev  = 0x00000001 // Open/Create/Delete device key
	diregDrv  = 0x00000002 // Open/Create/Delete driver key
	diregBoth = 0x00000004 // Delete both driver and Device key
)

// https://msdn.microsoft.com/it-it/library/windows/desktop/aa373931(v=vs.85).aspx
type guid struct {
	data1 uint32
	data2 uint16
	data3 uint16
	data4 [8]byte
}

func (g guid) String() string {
	return fmt.Sprintf("%08x-%04x-%04x-%02x%02x-%02x%02x%02x%02x%02x%02x",
		g.data1, g.data2, g.data3,
		g.data4[0], g.data4[1], g.data4[2], g.data4[3],
		g.data4[4], g.data4[5], g.data4[6], g.data4[7])
}

func classGuidsFromName(className string) ([]guid, error) {
	// Determine the number of GUIDs for className
	n := uint32(0)
	if err := setupDiClassGuidsFromNameInternal(className, nil, 0, &n); err != nil {
		// ignore error: UIDs array size too small
	}

	res := make([]guid, n)
	err := setupDiClassGuidsFromNameInternal(className, &res[0], n, &n)
	return res, err
}

const (
	digcfDefault         = 0x00000001 // only valid with digcfDeviceInterface
	digcfPresent         = 0x00000002
	digcfAllClasses      = 0x00000004
	digcfProfile         = 0x00000008
	digcfDeviceInterface = 0x00000010
)

type devicesSet syscall.Handle

func (g *guid) getDevicesSet() (devicesSet, error) {
	return setupDiGetClassDevs(g, nil, 0, digcfPresent)
}

func (set devicesSet) destroy() {
	setupDiDestroyDeviceInfoList(set)
}

// https://msdn.microsoft.com/en-us/library/windows/hardware/ff552344(v=vs.85).aspx
type devInfoData struct {
	size     uint32
	guid     guid
	devInst  uint32
	reserved uintptr
}

type deviceInfo struct {
	set  devicesSet
	data devInfoData
}

func (set devicesSet) getDeviceInfo(index int) (*deviceInfo, error) {
	result := &deviceInfo{set: set}

	result.data.size = uint32(unsafe.Sizeof(result.data))
	err := setupDiEnumDeviceInfo(set, uint32(index), &result.data)
	return result, err
}

func (dev *deviceInfo) getInstanceID() (string, error) {
	n := uint32(0)
	setupDiGetDeviceInstanceId(dev.set, &dev.data, nil, 0, &n)
	buff := make([]uint16, n)
	if err := setupDiGetDeviceInstanceId(dev.set, &dev.data, unsafe.Pointer(&buff[0]), uint32(len(buff)), &n); err != nil {
		return "", err
	}
	return windows.UTF16ToString(buff[:]), nil
}

func (dev *deviceInfo) openDevRegKey(scope dicsScope, hwProfile uint32, keyType uint32, samDesired regsam) (syscall.Handle, error) {
	return setupDiOpenDevRegKey(dev.set, &dev.data, scope, hwProfile, keyType, samDesired)
}

func nativeGetDetailedPortsList() ([]*PortDetails, error) {
	guids, err := classGuidsFromName("Ports")
	if err != nil {
		return nil, &PortEnumerationError{causedBy: err}
	}

	var res []*PortDetails
	for _, g := range guids {
		devsSet, err := g.getDevicesSet()
		if err != nil {
			return nil, &PortEnumerationError{causedBy: err}
		}
		defer devsSet.destroy()

		for i := 0; ; i++ {
			device, err := devsSet.getDeviceInfo(i)
			if err != nil {
				break
			}
			details := &PortDetails{}
			portName, err := retrievePortNameFromDevInfo(device)
			if err != nil {
				continue
			}
			if len(portName) < 3 || portName[0:3] != "COM" {
				// Accept only COM ports
				continue
			}
			details.Name = portName

			if err := retrievePortDetailsFromDevInfo(device, details); err != nil {
				return nil, &PortEnumerationError{causedBy: err}
			}
			res = append(res, details)
		}
	}
	return res, nil
}

func retrievePortNameFromDevInfo(device *deviceInfo) (string, error) {
	h, err := device.openDevRegKey(dicsFlagGlobal, 0, diregDev, keyRead)
	if err != nil {
		return "", err
	}
	defer syscall.RegCloseKey(h)

	var name [1024]uint16
	nameP := (*byte)(unsafe.Pointer(&name[0]))
	nameSize := uint32(len(name) * 2)
	if err := syscall.RegQueryValueEx(h, syscall.StringToUTF16Ptr("PortName"), nil, nil, nameP, &nameSize); err != nil {
		return "", err
	}
	return syscall.UTF16ToString(name[:]), nil
}

func retrievePortDetailsFromDevInfo(device *deviceInfo, details *PortDetails) error {
	deviceID, err := device.getInstanceID()
	if err != nil {
		return err
	}
	parseDeviceID(deviceID, details)

	/*	spdrpDeviceDesc returns a generic name, e.g.: "CDC-ACM", which will be the same for 2 identical devices attached
		while spdrpFriendlyName returns a specific name, e.g.: "CDC-ACM (COM44)",
		the result of spdrpFriendlyName is therefore unique and suitable as an alternative string to for a port choice */
	n := uint32(0)
	setupDiGetDeviceRegistryProperty(device.set, &device.data, /* spdrpDeviceDesc */ spdrpFriendlyName, nil, nil, 0, &n)
	buff := make([]uint16, n*2)
	buffP := (*byte)(unsafe.Pointer(&buff[0]))
	if setupDiGetDeviceRegistryProperty(device.set, &device.data, /* spdrpDeviceDesc */ spdrpFriendlyName, nil, buffP, n, &n) {
		details.Product = syscall.UTF16ToString(buff[:])
	}

	return nil
}
<|endoftext|>