<commit_before>package helper

import (
	"bytes"
	"crypto/x509"
	"encoding/asn1"
	"encoding/base64"
	"encoding/pem"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"sync"
	"time"

	"golang.org/x/crypto/ocsp"
)

var (
	method             = flag.String("method", "GET", "Method to use for fetching OCSP")
	urlOverride        = flag.String("url", "", "URL of OCSP responder to override")
	hostOverride       = flag.String("host", "", "Host header to override in HTTP request")
	tooSoon            = flag.Int("too-soon", 76, "If NextUpdate is fewer than this many hours in future, warn.")
	ignoreExpiredCerts = flag.Bool("ignore-expired-certs", false, "If a cert is expired, don't bother requesting OCSP.")
	expectStatus       = flag.Int("expect-status", 0, "Expect response to have this numeric status (0=Good, 1=Revoked, 2=Unknown); or -1 for no enforcement.")
	expectReason       = flag.Int("expect-reason", -1, "Expect response to have this numeric revocation reason (0=Unspecified, 1=KeyCompromise, etc); or -1 for no enforcement.")
)

// Config contains fields which control various behaviors of the
// checker's behavior.
type Config struct {
	method             string
	urlOverride        string
	hostOverride       string
	tooSoon            int
	ignoreExpiredCerts bool
	expectStatus       int
	expectReason       int
}

// DefaultConfig is a Config populated with the same defaults as if no
// command-line had been provided, so all retain their default value.
var DefaultConfig = Config{
	method:             *method,
	urlOverride:        *urlOverride,
	hostOverride:       *hostOverride,
	tooSoon:            *tooSoon,
	ignoreExpiredCerts: *ignoreExpiredCerts,
	expectStatus:       *expectStatus,
	expectReason:       *expectReason,
}

var parseFlagsOnce sync.Once

// ConfigFromFlags returns a Config whose values are populated from
// any command line flags passed by the user, or default values if not passed.
func ConfigFromFlags() Config {
	parseFlagsOnce.Do(func() {
		flag.Parse()
	})
	return Config{
		method:             *method,
		urlOverride:        *urlOverride,
		hostOverride:       *hostOverride,
		tooSoon:            *tooSoon,
		ignoreExpiredCerts: *ignoreExpiredCerts,
		expectStatus:       *expectStatus,
		expectReason:       *expectReason,
	}
}

// WithExpectStatus returns a new Config with the given expectStatus,
// and all other fields the same as the receiver.
func (template Config) WithExpectStatus(status int) Config {
	ret := template
	ret.expectStatus = status
	return ret
}

func getIssuer(cert *x509.Certificate) (*x509.Certificate, error) {
	if cert == nil {
		return nil, fmt.Errorf("nil certificate")
	}
	if len(cert.IssuingCertificateURL) == 0 {
		return nil, fmt.Errorf("No AIA information available, can't get issuer")
	}
	issuerURL := cert.IssuingCertificateURL[0]
	resp, err := http.Get(issuerURL)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	var issuer *x509.Certificate
	contentType := resp.Header.Get("Content-Type")
	if contentType == "application/x-pkcs7-mime" || contentType == "application/pkcs7-mime" {
		issuer, err = parseCMS(body)
	} else {
		issuer, err = parse(body)
	}
	if err != nil {
		return nil, fmt.Errorf("from %s: %s", issuerURL, err)
	}
	return issuer, nil
}

func parse(body []byte) (*x509.Certificate, error) {
	block, _ := pem.Decode(body)
	var der []byte
	if block == nil {
		der = body
	} else {
		der = block.Bytes
	}
	cert, err := x509.ParseCertificate(der)
	if err != nil {
		return nil, err
	}
	return cert, nil
}

// parseCMS parses certificates from CMS messages of type SignedData.
func parseCMS(body []byte) (*x509.Certificate, error) {
	type signedData struct {
		Version          int
		Digests          asn1.RawValue
		EncapContentInfo asn1.RawValue
		Certificates     asn1.RawValue
	}
	type cms struct {
		ContentType asn1.ObjectIdentifier
		SignedData  signedData `asn1:"explicit,tag:0"`
	}
	var msg cms
	_, err := asn1.Unmarshal(body, &msg)
	if err != nil {
		return nil, fmt.Errorf("parsing CMS: %s", err)
	}
	cert, err := x509.ParseCertificate(msg.SignedData.Certificates.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parsing CMS: %s", err)
	}
	return cert, nil
}

// Req makes an OCSP request using the given config for the PEM certificate in
// fileName, and returns the response.
func Req(fileName string, config Config) (*ocsp.Response, error) {
	contents, err := ioutil.ReadFile(fileName)
	if err != nil {
		return nil, err
	}
	return ReqDER(contents, config)
}

// ReqDER makes an OCSP request using the given config for the given DER-encoded
// certificate, and returns the response.
func ReqDER(der []byte, config Config) (*ocsp.Response, error) {
	cert, err := parse(der)
	if err != nil {
		return nil, fmt.Errorf("parsing certificate: %s", err)
	}
	if time.Now().After(cert.NotAfter) {
		if config.ignoreExpiredCerts {
			return nil, nil
		} else {
			return nil, fmt.Errorf("certificate expired %s ago: %s",
				time.Since(cert.NotAfter), cert.NotAfter)
		}
	}

	issuer, err := getIssuer(cert)
	if err != nil {
		return nil, fmt.Errorf("getting issuer: %s", err)
	}
	req, err := ocsp.CreateRequest(cert, issuer, nil)
	if err != nil {
		return nil, fmt.Errorf("creating OCSP request: %s", err)
	}

	ocspURL, err := getOCSPURL(cert, config.urlOverride)
	if err != nil {
		return nil, err
	}

	httpResp, err := sendHTTPRequest(req, ocspURL, config.method, config.hostOverride)
	if err != nil {
		return nil, err
	}
	fmt.Printf("HTTP %d\n", httpResp.StatusCode)
	for k, v := range httpResp.Header {
		for _, vv := range v {
			fmt.Printf("%s: %s\n", k, vv)
		}
	}
	if httpResp.StatusCode != 200 {
		return nil, fmt.Errorf("http status code %d", httpResp.StatusCode)
	}
	respBytes, err := ioutil.ReadAll(httpResp.Body)
	defer httpResp.Body.Close()
	if err != nil {
		return nil, err
	}
	if len(respBytes) == 0 {
		return nil, fmt.Errorf("empty response body")
	}
	return parseAndPrint(respBytes, cert, issuer, config)
}

func sendHTTPRequest(req []byte, ocspURL *url.URL, method string, host string) (*http.Response, error) {
	encodedReq := base64.StdEncoding.EncodeToString(req)
	var httpRequest *http.Request
	var err error
	if method == "GET" {
		ocspURL.Path = encodedReq
		fmt.Printf("Fetching %s\n", ocspURL.String())
		httpRequest, err = http.NewRequest("GET", ocspURL.String(), http.NoBody)
	} else if method == "POST" {
		fmt.Printf("POSTing request, reproduce with: curl -i --data-binary @- %s < <(base64 -d <<<%s)\n",
			ocspURL, encodedReq)
		httpRequest, err = http.NewRequest("POST", ocspURL.String(), bytes.NewBuffer(req))
	} else {
		return nil, fmt.Errorf("invalid method %s, expected GET or POST", method)
	}
	if err != nil {
		return nil, err
	}
	httpRequest.Header.Add("Content-Type", "application/ocsp-request")
	if host != "" {
		httpRequest.Host = host
	}
	client := http.Client{
		Timeout: 5 * time.Second,
	}

	return client.Do(httpRequest)
}

func getOCSPURL(cert *x509.Certificate, urlOverride string) (*url.URL, error) {
	var ocspServer string
	if urlOverride != "" {
		ocspServer = urlOverride
	} else if len(cert.OCSPServer) > 0 {
		ocspServer = cert.OCSPServer[0]
	} else {
		return nil, fmt.Errorf("no ocsp servers in cert")
	}
	ocspURL, err := url.Parse(ocspServer)
	if err != nil {
		return nil, fmt.Errorf("parsing URL: %s", err)
	}
	return ocspURL, nil
}

// checkSignerTimes checks that the OCSP response is within the
// validity window of whichever certificate signed it, and that that
// certificate is currently valid.
func checkSignerTimes(resp *ocsp.Response, issuer *x509.Certificate) error {
	var ocspSigner = issuer
	if delegatedSigner := resp.Certificate; delegatedSigner != nil {
		ocspSigner = delegatedSigner

		fmt.Printf("Using delegated OCSP signer from response: %s\n",
			base64.StdEncoding.EncodeToString(ocspSigner.Raw))
	}

	if resp.NextUpdate.After(ocspSigner.NotAfter) {
		return fmt.Errorf("OCSP response is valid longer than OCSP signer (%s): %s is after %s",
			ocspSigner.Subject, resp.NextUpdate, ocspSigner.NotAfter)
	}
	if resp.ThisUpdate.Before(ocspSigner.NotBefore) {
		return fmt.Errorf("OCSP response's validity begins before the OCSP signer's (%s): %s is before %s",
			ocspSigner.Subject, resp.ThisUpdate, ocspSigner.NotBefore)
	}

	if time.Now().After(ocspSigner.NotAfter) {
		return fmt.Errorf("OCSP signer (%s) expired at %s", ocspSigner.Subject, ocspSigner.NotAfter)
	}
	if time.Now().Before(ocspSigner.NotBefore) {
		return fmt.Errorf("OCSP signer (%s) not valid until %s", ocspSigner.Subject, ocspSigner.NotBefore)
	}
	return nil
}

func parseAndPrint(respBytes []byte, cert, issuer *x509.Certificate, config Config) (*ocsp.Response, error) {
	fmt.Printf("\nDecoding body: %s\n", base64.StdEncoding.EncodeToString(respBytes))
	resp, err := ocsp.ParseResponseForCert(respBytes, cert, issuer)
	if err != nil {
		return nil, fmt.Errorf("parsing response: %s", err)
	}
	if config.expectStatus != -1 && resp.Status != config.expectStatus {
		return nil, fmt.Errorf("wrong CertStatus %d, expected %d", resp.Status, config.expectStatus)
	}
	if config.expectReason != -1 && resp.RevocationReason != config.expectReason {
		return nil, fmt.Errorf("wrong RevocationReason %d, expected %d", resp.RevocationReason, config.expectReason)
	}
	timeTilExpiry := time.Until(resp.NextUpdate)
	tooSoonDuration := time.Duration(config.tooSoon) * time.Hour
	if timeTilExpiry < tooSoonDuration {
		return nil, fmt.Errorf("NextUpdate is too soon: %s", timeTilExpiry)
	}

	err = checkSignerTimes(resp, issuer)
	if err != nil {
		return nil, fmt.Errorf("checking signature on delegated signer: %s", err)
	}

	fmt.Printf("\n")
	fmt.Printf("Good response:\n")
	fmt.Printf("  CertStatus %d\n", resp.Status)
	fmt.Printf("  SerialNumber %036x\n", resp.SerialNumber)
	fmt.Printf("  ProducedAt %s\n", resp.ProducedAt)
	fmt.Printf("  ThisUpdate %s\n", resp.ThisUpdate)
	fmt.Printf("  NextUpdate %s\n", resp.NextUpdate)
	fmt.Printf("  RevokedAt %s\n", resp.RevokedAt)
	fmt.Printf("  RevocationReason %d\n", resp.RevocationReason)
	fmt.Printf("  SignatureAlgorithm %s\n", resp.SignatureAlgorithm)
	fmt.Printf("  Extensions %#v\n", resp.Extensions)
	if resp.Certificate == nil {
		fmt.Printf("  Certificate: nil\n")
	} else {
		fmt.Print("  Certificate:\n")
		fmt.Printf("    Subject: %s\n", resp.Certificate.Subject)
		fmt.Printf("    Issuer: %s\n", resp.Certificate.Issuer)
		fmt.Printf("    NotBefore: %s\n", resp.Certificate.NotBefore)
		fmt.Printf("    NotAfter: %s\n", resp.Certificate.NotAfter)
	}
	return resp, nil
}
<commit_msg>Update default value of ocspchecker -expectStatus (#4922)<commit_after>package helper

import (
	"bytes"
	"crypto/x509"
	"encoding/asn1"
	"encoding/base64"
	"encoding/pem"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"sync"
	"time"

	"golang.org/x/crypto/ocsp"
)

var (
	method             = flag.String("method", "GET", "Method to use for fetching OCSP")
	urlOverride        = flag.String("url", "", "URL of OCSP responder to override")
	hostOverride       = flag.String("host", "", "Host header to override in HTTP request")
	tooSoon            = flag.Int("too-soon", 76, "If NextUpdate is fewer than this many hours in future, warn.")
	ignoreExpiredCerts = flag.Bool("ignore-expired-certs", false, "If a cert is expired, don't bother requesting OCSP.")
	expectStatus       = flag.Int("expect-status", -1, "Expect response to have this numeric status (0=Good, 1=Revoked, 2=Unknown); or -1 for no enforcement.")
	expectReason       = flag.Int("expect-reason", -1, "Expect response to have this numeric revocation reason (0=Unspecified, 1=KeyCompromise, etc); or -1 for no enforcement.")
)

// Config contains fields which control various behaviors of the
// checker's behavior.
type Config struct {
	method             string
	urlOverride        string
	hostOverride       string
	tooSoon            int
	ignoreExpiredCerts bool
	expectStatus       int
	expectReason       int
}

// DefaultConfig is a Config populated with the same defaults as if no
// command-line had been provided, so all retain their default value.
var DefaultConfig = Config{
	method:             *method,
	urlOverride:        *urlOverride,
	hostOverride:       *hostOverride,
	tooSoon:            *tooSoon,
	ignoreExpiredCerts: *ignoreExpiredCerts,
	expectStatus:       *expectStatus,
	expectReason:       *expectReason,
}

var parseFlagsOnce sync.Once

// ConfigFromFlags returns a Config whose values are populated from
// any command line flags passed by the user, or default values if not passed.
func ConfigFromFlags() Config {
	parseFlagsOnce.Do(func() {
		flag.Parse()
	})
	return Config{
		method:             *method,
		urlOverride:        *urlOverride,
		hostOverride:       *hostOverride,
		tooSoon:            *tooSoon,
		ignoreExpiredCerts: *ignoreExpiredCerts,
		expectStatus:       *expectStatus,
		expectReason:       *expectReason,
	}
}

// WithExpectStatus returns a new Config with the given expectStatus,
// and all other fields the same as the receiver.
func (template Config) WithExpectStatus(status int) Config {
	ret := template
	ret.expectStatus = status
	return ret
}

func getIssuer(cert *x509.Certificate) (*x509.Certificate, error) {
	if cert == nil {
		return nil, fmt.Errorf("nil certificate")
	}
	if len(cert.IssuingCertificateURL) == 0 {
		return nil, fmt.Errorf("No AIA information available, can't get issuer")
	}
	issuerURL := cert.IssuingCertificateURL[0]
	resp, err := http.Get(issuerURL)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	var issuer *x509.Certificate
	contentType := resp.Header.Get("Content-Type")
	if contentType == "application/x-pkcs7-mime" || contentType == "application/pkcs7-mime" {
		issuer, err = parseCMS(body)
	} else {
		issuer, err = parse(body)
	}
	if err != nil {
		return nil, fmt.Errorf("from %s: %s", issuerURL, err)
	}
	return issuer, nil
}

func parse(body []byte) (*x509.Certificate, error) {
	block, _ := pem.Decode(body)
	var der []byte
	if block == nil {
		der = body
	} else {
		der = block.Bytes
	}
	cert, err := x509.ParseCertificate(der)
	if err != nil {
		return nil, err
	}
	return cert, nil
}

// parseCMS parses certificates from CMS messages of type SignedData.
func parseCMS(body []byte) (*x509.Certificate, error) {
	type signedData struct {
		Version          int
		Digests          asn1.RawValue
		EncapContentInfo asn1.RawValue
		Certificates     asn1.RawValue
	}
	type cms struct {
		ContentType asn1.ObjectIdentifier
		SignedData  signedData `asn1:"explicit,tag:0"`
	}
	var msg cms
	_, err := asn1.Unmarshal(body, &msg)
	if err != nil {
		return nil, fmt.Errorf("parsing CMS: %s", err)
	}
	cert, err := x509.ParseCertificate(msg.SignedData.Certificates.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parsing CMS: %s", err)
	}
	return cert, nil
}

// Req makes an OCSP request using the given config for the PEM certificate in
// fileName, and returns the response.
func Req(fileName string, config Config) (*ocsp.Response, error) {
	contents, err := ioutil.ReadFile(fileName)
	if err != nil {
		return nil, err
	}
	return ReqDER(contents, config)
}

// ReqDER makes an OCSP request using the given config for the given DER-encoded
// certificate, and returns the response.
func ReqDER(der []byte, config Config) (*ocsp.Response, error) {
	cert, err := parse(der)
	if err != nil {
		return nil, fmt.Errorf("parsing certificate: %s", err)
	}
	if time.Now().After(cert.NotAfter) {
		if config.ignoreExpiredCerts {
			return nil, nil
		} else {
			return nil, fmt.Errorf("certificate expired %s ago: %s",
				time.Since(cert.NotAfter), cert.NotAfter)
		}
	}

	issuer, err := getIssuer(cert)
	if err != nil {
		return nil, fmt.Errorf("getting issuer: %s", err)
	}
	req, err := ocsp.CreateRequest(cert, issuer, nil)
	if err != nil {
		return nil, fmt.Errorf("creating OCSP request: %s", err)
	}

	ocspURL, err := getOCSPURL(cert, config.urlOverride)
	if err != nil {
		return nil, err
	}

	httpResp, err := sendHTTPRequest(req, ocspURL, config.method, config.hostOverride)
	if err != nil {
		return nil, err
	}
	fmt.Printf("HTTP %d\n", httpResp.StatusCode)
	for k, v := range httpResp.Header {
		for _, vv := range v {
			fmt.Printf("%s: %s\n", k, vv)
		}
	}
	if httpResp.StatusCode != 200 {
		return nil, fmt.Errorf("http status code %d", httpResp.StatusCode)
	}
	respBytes, err := ioutil.ReadAll(httpResp.Body)
	defer httpResp.Body.Close()
	if err != nil {
		return nil, err
	}
	if len(respBytes) == 0 {
		return nil, fmt.Errorf("empty response body")
	}
	return parseAndPrint(respBytes, cert, issuer, config)
}

func sendHTTPRequest(req []byte, ocspURL *url.URL, method string, host string) (*http.Response, error) {
	encodedReq := base64.StdEncoding.EncodeToString(req)
	var httpRequest *http.Request
	var err error
	if method == "GET" {
		ocspURL.Path = encodedReq
		fmt.Printf("Fetching %s\n", ocspURL.String())
		httpRequest, err = http.NewRequest("GET", ocspURL.String(), http.NoBody)
	} else if method == "POST" {
		fmt.Printf("POSTing request, reproduce with: curl -i --data-binary @- %s < <(base64 -d <<<%s)\n",
			ocspURL, encodedReq)
		httpRequest, err = http.NewRequest("POST", ocspURL.String(), bytes.NewBuffer(req))
	} else {
		return nil, fmt.Errorf("invalid method %s, expected GET or POST", method)
	}
	if err != nil {
		return nil, err
	}
	httpRequest.Header.Add("Content-Type", "application/ocsp-request")
	if host != "" {
		httpRequest.Host = host
	}
	client := http.Client{
		Timeout: 5 * time.Second,
	}

	return client.Do(httpRequest)
}

func getOCSPURL(cert *x509.Certificate, urlOverride string) (*url.URL, error) {
	var ocspServer string
	if urlOverride != "" {
		ocspServer = urlOverride
	} else if len(cert.OCSPServer) > 0 {
		ocspServer = cert.OCSPServer[0]
	} else {
		return nil, fmt.Errorf("no ocsp servers in cert")
	}
	ocspURL, err := url.Parse(ocspServer)
	if err != nil {
		return nil, fmt.Errorf("parsing URL: %s", err)
	}
	return ocspURL, nil
}

// checkSignerTimes checks that the OCSP response is within the
// validity window of whichever certificate signed it, and that that
// certificate is currently valid.
func checkSignerTimes(resp *ocsp.Response, issuer *x509.Certificate) error {
	var ocspSigner = issuer
	if delegatedSigner := resp.Certificate; delegatedSigner != nil {
		ocspSigner = delegatedSigner

		fmt.Printf("Using delegated OCSP signer from response: %s\n",
			base64.StdEncoding.EncodeToString(ocspSigner.Raw))
	}

	if resp.NextUpdate.After(ocspSigner.NotAfter) {
		return fmt.Errorf("OCSP response is valid longer than OCSP signer (%s): %s is after %s",
			ocspSigner.Subject, resp.NextUpdate, ocspSigner.NotAfter)
	}
	if resp.ThisUpdate.Before(ocspSigner.NotBefore) {
		return fmt.Errorf("OCSP response's validity begins before the OCSP signer's (%s): %s is before %s",
			ocspSigner.Subject, resp.ThisUpdate, ocspSigner.NotBefore)
	}

	if time.Now().After(ocspSigner.NotAfter) {
		return fmt.Errorf("OCSP signer (%s) expired at %s", ocspSigner.Subject, ocspSigner.NotAfter)
	}
	if time.Now().Before(ocspSigner.NotBefore) {
		return fmt.Errorf("OCSP signer (%s) not valid until %s", ocspSigner.Subject, ocspSigner.NotBefore)
	}
	return nil
}

func parseAndPrint(respBytes []byte, cert, issuer *x509.Certificate, config Config) (*ocsp.Response, error) {
	fmt.Printf("\nDecoding body: %s\n", base64.StdEncoding.EncodeToString(respBytes))
	resp, err := ocsp.ParseResponseForCert(respBytes, cert, issuer)
	if err != nil {
		return nil, fmt.Errorf("parsing response: %s", err)
	}
	if config.expectStatus != -1 && resp.Status != config.expectStatus {
		return nil, fmt.Errorf("wrong CertStatus %d, expected %d", resp.Status, config.expectStatus)
	}
	if config.expectReason != -1 && resp.RevocationReason != config.expectReason {
		return nil, fmt.Errorf("wrong RevocationReason %d, expected %d", resp.RevocationReason, config.expectReason)
	}
	timeTilExpiry := time.Until(resp.NextUpdate)
	tooSoonDuration := time.Duration(config.tooSoon) * time.Hour
	if timeTilExpiry < tooSoonDuration {
		return nil, fmt.Errorf("NextUpdate is too soon: %s", timeTilExpiry)
	}

	err = checkSignerTimes(resp, issuer)
	if err != nil {
		return nil, fmt.Errorf("checking signature on delegated signer: %s", err)
	}

	fmt.Printf("\n")
	fmt.Printf("Good response:\n")
	fmt.Printf("  CertStatus %d\n", resp.Status)
	fmt.Printf("  SerialNumber %036x\n", resp.SerialNumber)
	fmt.Printf("  ProducedAt %s\n", resp.ProducedAt)
	fmt.Printf("  ThisUpdate %s\n", resp.ThisUpdate)
	fmt.Printf("  NextUpdate %s\n", resp.NextUpdate)
	fmt.Printf("  RevokedAt %s\n", resp.RevokedAt)
	fmt.Printf("  RevocationReason %d\n", resp.RevocationReason)
	fmt.Printf("  SignatureAlgorithm %s\n", resp.SignatureAlgorithm)
	fmt.Printf("  Extensions %#v\n", resp.Extensions)
	if resp.Certificate == nil {
		fmt.Printf("  Certificate: nil\n")
	} else {
		fmt.Print("  Certificate:\n")
		fmt.Printf("    Subject: %s\n", resp.Certificate.Subject)
		fmt.Printf("    Issuer: %s\n", resp.Certificate.Issuer)
		fmt.Printf("    NotBefore: %s\n", resp.Certificate.NotBefore)
		fmt.Printf("    NotAfter: %s\n", resp.Certificate.NotAfter)
	}
	return resp, nil
}
<|endoftext|>


<commit_before>// Automatically downloads and configures Steam grid images for all games in a
// given Steam installation.
package main

import (
	"errors"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/boppreh/go-ui"
)

// Prints an error and quits.
func errorAndExit(err error) {
	goui.Error("An unexpected error occurred:", err.Error())
	os.Exit(1)
}

func main() {
	goui.Start(func() {
		http.DefaultTransport.(*http.Transport).ResponseHeaderTimeout = time.Second * 10

		descriptions := make(chan string)
		progress := make(chan int)

		go goui.Progress("SteamGrid", descriptions, progress, func() { os.Exit(1) })

		startApplication(descriptions, progress)
	})
}

func startApplication(descriptions chan string, progress chan int) {
	descriptions <- "Loading overlays..."
	overlays, err := LoadOverlays(filepath.Join(filepath.Dir(os.Args[0]), "overlays by category"))
	if err != nil {
		errorAndExit(err)
	}
	if len(overlays) == 0 {
		// I'm trying to use a message box here, but for some reason the
		// message appears twice and there's an error a closed channel.
		fmt.Println("No overlays", "No category overlays found. You can put overlay images in the folder 'overlays by category', where the filename is the game category.\n\nContinuing without overlays...")
	}

	descriptions <- "Looking for Steam directory..."
	installationDir, err := GetSteamInstallation()
	if err != nil {
		errorAndExit(err)
	}

	descriptions <- "Loading users..."
	users, err := GetUsers(installationDir)
	if err != nil {
		errorAndExit(err)
	}
	if len(users) == 0 {
		errorAndExit(errors.New("No users found at Steam/userdata. Have you used Steam before in this computer?"))
	}

	nOverlaysApplied := 0
	nDownloaded := 0
	notFounds := make([]*Game, 0)
	searchFounds := make([]*Game, 0)

	for _, user := range users {
		descriptions <- "Loading games for " + user.Name

		RestoreBackup(user)

		games := GetGames(user)

		i := 0
		for _, game := range games {
			fmt.Println(game.Name)

			i++
			progress <- i * 100 / len(games)
			var name string
			if game.Name != "" {
				name = game.Name
			} else {
				name = "unknown game with id " + game.Id
			}
			descriptions <- fmt.Sprintf("Processing %v (%v/%v)",
				name, i, len(games))

			downloaded, found, fromSearch, err := DownloadImage(game, user)
			if err != nil {
				errorAndExit(err)
			}
			if downloaded {
				nDownloaded++
			}
			if !found {
				notFounds = append(notFounds, game)
				continue
			}
			if fromSearch {
				searchFounds = append(searchFounds, game)
			}

			err = BackupGame(game)
			if err != nil {
				errorAndExit(err)
			}

			applied, err := ApplyOverlay(game, overlays)
			if err != nil {
				errorAndExit(err)
			}
			if applied {
				nOverlaysApplied++
			}
		}
	}

	close(progress)

	message := fmt.Sprintf("%v images downloaded and %v overlays applied.\n\n", nDownloaded, nOverlaysApplied)
	if len(searchFounds) >= 1 {
		message += fmt.Sprintf("%v images were found with a Google search and may not be accurate:\n", len(searchFounds))
		for _, game := range searchFounds {
			message += fmt.Sprintf("* %v (steam id %v)\n", game.Name, game.Id)
		}

		message += "\n\n"
	}

	if len(notFounds) >= 1 {
		message += fmt.Sprintf("%v images could not be found anywhere:\n", len(notFounds))
		for _, game := range notFounds {
			message += fmt.Sprintf("- %v (id %v)\n", game.Name, game.Id)
		}

		message += "\n\n"
	}
	message += "Open Steam in grid view to see the results!"

	goui.Info("Results", message)
}
<commit_msg>Don't quit after finding an encoding error<commit_after>// Automatically downloads and configures Steam grid images for all games in a
// given Steam installation.
package main

import (
	"errors"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/boppreh/go-ui"
)

// Prints an error and quits.
func errorAndExit(err error) {
	goui.Error("An unexpected error occurred:", err.Error())
	os.Exit(1)
}

func main() {
	goui.Start(func() {
		http.DefaultTransport.(*http.Transport).ResponseHeaderTimeout = time.Second * 10

		descriptions := make(chan string)
		progress := make(chan int)

		go goui.Progress("SteamGrid", descriptions, progress, func() { os.Exit(1) })

		startApplication(descriptions, progress)
	})
}

func startApplication(descriptions chan string, progress chan int) {
	descriptions <- "Loading overlays..."
	overlays, err := LoadOverlays(filepath.Join(filepath.Dir(os.Args[0]), "overlays by category"))
	if err != nil {
		errorAndExit(err)
	}
	if len(overlays) == 0 {
		// I'm trying to use a message box here, but for some reason the
		// message appears twice and there's an error a closed channel.
		fmt.Println("No overlays", "No category overlays found. You can put overlay images in the folder 'overlays by category', where the filename is the game category.\n\nContinuing without overlays...")
	}

	descriptions <- "Looking for Steam directory..."
	installationDir, err := GetSteamInstallation()
	if err != nil {
		errorAndExit(err)
	}

	descriptions <- "Loading users..."
	users, err := GetUsers(installationDir)
	if err != nil {
		errorAndExit(err)
	}
	if len(users) == 0 {
		errorAndExit(errors.New("No users found at Steam/userdata. Have you used Steam before in this computer?"))
	}

	nOverlaysApplied := 0
	nDownloaded := 0
	notFounds := make([]*Game, 0)
	searchFounds := make([]*Game, 0)
	errors := make([]*Game, 0)
	errorMessages := make([]string, 0)

	for _, user := range users {
		descriptions <- "Loading games for " + user.Name

		RestoreBackup(user)

		games := GetGames(user)

		i := 0
		for _, game := range games {
			fmt.Println(game.Name)

			i++
			progress <- i * 100 / len(games)
			var name string
			if game.Name != "" {
				name = game.Name
			} else {
				name = "unknown game with id " + game.Id
			}
			descriptions <- fmt.Sprintf("Processing %v (%v/%v)",
				name, i, len(games))

			downloaded, found, fromSearch, err := DownloadImage(game, user)
			if err != nil {
				errorAndExit(err)
			}
			if downloaded {
				nDownloaded++
			}
			if !found {
				notFounds = append(notFounds, game)
				continue
			}
			if fromSearch {
				searchFounds = append(searchFounds, game)
			}

			err = BackupGame(game)
			if err != nil {
				errorAndExit(err)
			}

			applied, err := ApplyOverlay(game, overlays)
			if err != nil {
				print(err.Error(), "\n")
				errors = append(errors, game)
				errorMessages = append(errorMessages, err.Error())
			}
			if applied {
				nOverlaysApplied++
			}
		}
	}

	close(progress)

	message := fmt.Sprintf("%v images downloaded and %v overlays applied.\n\n", nDownloaded, nOverlaysApplied)
	if len(searchFounds) >= 1 {
		message += fmt.Sprintf("%v images were found with a Google search and may not be accurate:\n", len(searchFounds))
		for _, game := range searchFounds {
			message += fmt.Sprintf("* %v (steam id %v)\n", game.Name, game.Id)
		}

		message += "\n\n"
	}

	if len(notFounds) >= 1 {
		message += fmt.Sprintf("%v images could not be found anywhere:\n", len(notFounds))
		for _, game := range notFounds {
			message += fmt.Sprintf("- %v (id %v)\n", game.Name, game.Id)
		}

		message += "\n\n"
	}

	if len(errors) >= 1 {
		message += fmt.Sprintf("%v images were found but had errors and could not be overlaid:\n", len(errors))
		for i, game := range errors {
			message += fmt.Sprintf("- %v (id %v) (%v)\n", game.Name, game.Id, errorMessages[i])
		}

		message += "\n\n"
	}

	message += "Open Steam in grid view to see the results!"

	goui.Info("Results", message)
}
<|endoftext|>


<commit_before>package svfs

import (
	"regexp"
	"strings"

	"golang.org/x/net/context"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"

	"github.com/ncw/swift"
)

var SegmentRegex = regexp.MustCompile("^.+_segments$")

type Root struct {
	*Directory
}

func (r *Root) Create(ctx context.Context, req *fuse.CreateRequest, resp *fuse.CreateResponse) (fs.Node, fs.Handle, error) {
	// TODO : implement container creation
	return nil, nil, fuse.ENOTSUP
}

func (r *Root) Remove(ctx context.Context, req *fuse.RemoveRequest) error {
	// TODO : implement container removal
	return fuse.ENOTSUP
}

func (r *Root) Rename(ctx context.Context, req *fuse.RenameRequest, newDir fs.Node) error {
	return fuse.ENOTSUP
}

func (r *Root) ReadDirAll(ctx context.Context) (entries []fuse.Dirent, err error) {
	var (
		baseC = make(map[string]swift.Container)
		segC  = make(map[string]swift.Container)
	)

	if len(r.children) > 0 {
		for _, c := range r.children {
			entries = append(entries, c.Export())
		}
		return entries, nil
	}

	// Retrieve base container
	c, err := r.s.ContainersAll(nil)
	if err != nil {
		return nil, err
	}

	// Find segments
	for _, container := range c {
		name := container.Name
		if !SegmentRegex.Match([]byte(name)) {
			baseC[name] = container
			continue
		}
		if SegmentRegex.Match([]byte(name)) {
			segC[strings.TrimSuffix(name, "_segments")] = container
			continue
		}
	}

	// Register children
	for name, container := range baseC {
		segment := segC[name]
		child := Container{
			Directory: &Directory{
				s:    r.s,
				c:    &container,
				path: "",
				name: name,
			},
			cs: &segment,
		}

		r.children = append(r.children, &child)
		entries = append(entries, child.Export())
	}

	return entries, nil
}

var (
	_ Node           = (*Root)(nil)
	_ fs.Node        = (*Root)(nil)
	_ fs.NodeCreater = (*Root)(nil)
	_ fs.NodeRemover = (*Root)(nil)
	_ fs.NodeRenamer = (*Root)(nil)
)
<commit_msg>Bugfix: containers hold a common reference to the last container listed<commit_after>package svfs

import (
	"regexp"
	"strings"

	"golang.org/x/net/context"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"

	"github.com/ncw/swift"
)

var SegmentRegex = regexp.MustCompile("^.+_segments$")

type Root struct {
	*Directory
}

func (r *Root) Create(ctx context.Context, req *fuse.CreateRequest, resp *fuse.CreateResponse) (fs.Node, fs.Handle, error) {
	// TODO : implement container creation
	return nil, nil, fuse.ENOTSUP
}

func (r *Root) Remove(ctx context.Context, req *fuse.RemoveRequest) error {
	// TODO : implement container removal
	return fuse.ENOTSUP
}

func (r *Root) Rename(ctx context.Context, req *fuse.RenameRequest, newDir fs.Node) error {
	return fuse.ENOTSUP
}

func (r *Root) ReadDirAll(ctx context.Context) (entries []fuse.Dirent, err error) {
	var (
		baseC = make(map[string]*swift.Container)
		segC  = make(map[string]*swift.Container)
	)

	if len(r.children) > 0 {
		for _, c := range r.children {
			entries = append(entries, c.Export())
		}
		return entries, nil
	}

	// Retrieve base container
	cs, err := r.s.ContainersAll(nil)
	if err != nil {
		return nil, err
	}

	// Find segments
	for _, container := range cs {
		c := container
		name := container.Name
		if !SegmentRegex.Match([]byte(name)) {
			baseC[name] = &c
			continue
		}
		if SegmentRegex.Match([]byte(name)) {
			segC[strings.TrimSuffix(name, "_segments")] = &c
			continue
		}
	}

	// Register children
	for name, container := range baseC {
		c := container
		segment := segC[name]
		child := Container{
			Directory: &Directory{
				s:    r.s,
				c:    c,
				path: "",
				name: name,
			},
			cs: segment,
		}

		r.children = append(r.children, &child)
		entries = append(entries, child.Export())
	}

	return entries, nil
}

var (
	_ Node           = (*Root)(nil)
	_ fs.Node        = (*Root)(nil)
	_ fs.NodeCreater = (*Root)(nil)
	_ fs.NodeRemover = (*Root)(nil)
	_ fs.NodeRenamer = (*Root)(nil)
)
<|endoftext|>


<commit_before>package util

import (
	"database/sql"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
)

func DBFromConfigDatabase(database string) (*sql.DB, error) {
	db, err := sql.Open("mysql", database+"?parseTime=true")
	if err != nil {
		return nil, fmt.Errorf("failed to create database with database(%v): %v", database, err)
	}
	if err := db.Ping(); err != nil {
		return nil, err
	}
	return db, nil
}
<commit_msg>veyron/services/identity: Remove unused sql_config file.<commit_after><|endoftext|>


<commit_before>package reporter

import (
  "os"
  "path"
  "runtime"
)

// DefaultConfigFile returns the path to Hearthstone's logging config file.
//
// It returns the expected file path, assuming a standard game installation.
func DefaultConfigFile() string {
  var configRoot string
  if runtime.GOOS == "windows" {
    configRoot = os.Getenv("LOCALAPPDATA")
  } else {
    // OSX.
    configRoot = path.Join(os.Getenv("HOME"), "Library", "Preferences")
  }
  return path.Join(configRoot, "Blizzard", "Hearthstone", "log.config")
}

// DefaultLogPath returns the path to Hearthstone's logging output file.
//
// It returns the expected file path, assuming a standard game installation.
func DefaultLogFile() string {
  if runtime.GOOS == "windows" {
    var programFiles string
    if runtime.GOARCH == "amd64" {
      programFiles = "C:\\Program Files (x86)"
    } else {
      programFiles = "C:\\Program Files"
    }
    return path.Join(programFiles, "Hearthstone", "Hearthstone_data",
        "output_log.txt")
  }

  // OSX.
  return path.Join(os.Getenv("HOME"), "Library", "Logs", "Unity", "Player.log")
}
<commit_msg>Windows path fix, take 4.<commit_after>package reporter

import (
  "os"
  "path"
  "runtime"
)

// DefaultConfigFile returns the path to Hearthstone's logging config file.
//
// It returns the expected file path, assuming a standard game installation.
func DefaultConfigFile() string {
  var configRoot string
  if runtime.GOOS == "windows" {
    configRoot = os.Getenv("LOCALAPPDATA")
  } else {
    // OSX.
    configRoot = path.Join(os.Getenv("HOME"), "Library", "Preferences")
  }
  return path.Join(configRoot, "Blizzard", "Hearthstone", "log.config")
}

// DefaultLogPath returns the path to Hearthstone's logging output file.
//
// It returns the expected file path, assuming a standard game installation.
func DefaultLogFile() string {
  if runtime.GOOS == "windows" {
    var programFiles string
    if runtime.GOARCH == "amd64" {
      programFiles = path.Join("C:", "Program Files (x86)")
    } else {
      programFiles = path.Join("C:", "Program Files")
    }
    return path.Join(programFiles, "Hearthstone", "Hearthstone_data",
        "output_log.txt")
  }

  // OSX.
  return path.Join(os.Getenv("HOME"), "Library", "Logs", "Unity", "Player.log")
}
<|endoftext|>


<commit_before>package api

// ecosystem_helpers_test.go has a bunch of helper functions to make setting up
// large ecosystem tests easier.
//
// List of helper functions:
//    addStorageToAllHosts // adds a storage folder to every host
//    announceAllHosts     // announce all hosts to the network (and mine a block)
//    fullyConnectNodes    // connects each server tester to all the others
//    fundAllNodes         // mines blocks until all server testers have money
//    synchronizationCheck // checks that all server testers have the same recent block
//    waitForBlock         // block until the provided block is the most recent block for all server testers

import (
	"errors"
	"fmt"
	"net/url"
	"time"

	"github.com/NebulousLabs/Sia/types"
)

// addStorageToAllHosts adds a storage folder with a bunch of storage to each
// host.
func addStorageToAllHosts(sts []*serverTester) error {
	for _, st := range sts {
		values := url.Values{}
		values.Set("path", st.dir)
		values.Set("size", "1048576")
		err := st.stdPostAPI("/host/storage/folders/add", values)
		if err != nil {
			return err
		}
	}
	return nil
}

// announceAllHosts will announce every host in the tester set to the
// blockchain.
func announceAllHosts(sts []*serverTester) error {
	// Check that all announcements will be on the same chain.
	_, err := synchronizationCheck(sts)
	if err != nil {
		return err
	}

	// Grab the initial transaction pool size to know how many total transactions
	// there should be after announcement.
	initialTpoolSize := len(sts[0].tpool.TransactionList())

	// Announce each host.
	for _, st := range sts {
		// Set the host to be accepting contracts.
		acceptingContractsValues := url.Values{}
		acceptingContractsValues.Set("acceptingcontracts", "true")
		err = st.stdPostAPI("/host", acceptingContractsValues)
		if err != nil {
			return err
		}

		// Fetch the host net address.
		var hg HostGET
		err = st.getAPI("/host", &hg)
		if err != nil {
			return err
		}

		// Make the announcement.
		announceValues := url.Values{}
		announceValues.Set("address", string(hg.ExternalSettings.NetAddress))
		err = st.stdPostAPI("/host/announce", announceValues)
		if err != nil {
			return err
		}
	}

	// Wait until all of the transactions have propagated to all of the nodes.
	//
	// TODO: Replace this direct transaction pool call with a call to the
	// /transactionpool endpoint.
	//
	// TODO: At some point the number of transactions needed to make an
	// announcement may change. Currently its 2.
	for i := 0; i < 50; i++ {
		if len(sts[0].tpool.TransactionList()) == len(sts)*2+initialTpoolSize {
			break
		}
		time.Sleep(time.Millisecond * 100)
	}
	if len(sts[0].tpool.TransactionList()) < len(sts)*2+initialTpoolSize {
		return fmt.Errorf("Host announcements do not seem to have propagated to the leader's tpool: %v, %v, %v", len(sts), len(sts[0].tpool.TransactionList())+initialTpoolSize, initialTpoolSize)
	}

	// Mine a block and then wait for all of the nodes to syncrhonize to it.
	_, err = sts[0].miner.AddBlock()
	if err != nil {
		return err
	}
	// Block until the block propagated to all nodes
	for _, st := range sts[1:] {
		err = waitForBlock(sts[0].cs.CurrentBlock().ID(), st)
		if err != nil {
			return (err)
		}
	}
	// Check if all nodes are on the same block now
	_, err = synchronizationCheck(sts)
	if err != nil {
		return err
	}

	// Block until every node has completed the scan of every other node, so
	// that each node has a full hostdb.
	for _, st := range sts {
		var ah HostdbActiveGET
		for i := 0; i < 100; i++ {
			err = st.getAPI("/hostdb/active", &ah)
			if err != nil {
				return err
			}
			if len(ah.Hosts) >= len(sts) {
				break
			}
			time.Sleep(time.Millisecond * 100)
		}
		if len(ah.Hosts) < len(sts) {
			return errors.New("one of the nodes hostdbs was unable to find at least one host announcement")
		}
	}
	return nil
}

// fullyConnectNodes takes a bunch of tester nodes and connects each to the
// other, creating a fully connected graph so that everyone is on the same
// chain.
//
// After connecting the nodes, it verifies that all the nodes have
// synchronized.
func fullyConnectNodes(sts []*serverTester) error {
	for i, sta := range sts {
		var gg GatewayGET
		err := sta.getAPI("/gateway", &gg)
		if err != nil {
			return err
		}

		// Connect this node to every other node.
		for _, stb := range sts[i+1:] {
			// Try connecting to the other node until both have the other in
			// their peer list.
			err = retry(100, time.Millisecond*100, func() error {
				// NOTE: this check depends on string-matching an error in the
				// gateway. If that error changes at all, this string will need to
				// be updated.
				err := stb.stdPostAPI("/gateway/connect/"+string(gg.NetAddress), nil)
				if err != nil && err.Error() != "already connected to this peer" {
					return err
				}

				// Check that the gateways are connected.
				bToA := false
				aToB := false
				var ggb GatewayGET
				err = stb.getAPI("/gateway", &ggb)
				if err != nil {
					return err
				}
				for _, peer := range ggb.Peers {
					if peer.NetAddress == gg.NetAddress {
						bToA = true
						break
					}
				}
				err = sta.getAPI("/gateway", &gg)
				if err != nil {
					return err
				}
				for _, peer := range gg.Peers {
					if peer.NetAddress == ggb.NetAddress {
						aToB = true
						break
					}
				}
				if !aToB || !bToA {
					return fmt.Errorf("called connect between two nodes, but they are not peers: %v %v %v %v %v %v", aToB, bToA, gg.NetAddress, ggb.NetAddress, gg.Peers, ggb.Peers)
				}
				return nil

			})
			if err != nil {
				return err
			}
		}
	}

	// Perform a synchronization check.
	_, err := synchronizationCheck(sts)
	return err
}

// fundAllNodes will make sure that each node has mined a block in the longest
// chain, then will mine enough blocks that the miner payouts manifest in the
// wallets of each node.
func fundAllNodes(sts []*serverTester) error {
	// Check that all of the nodes are synchronized.
	chainTip, err := synchronizationCheck(sts)
	if err != nil {
		return err
	}

	// Mine a block for each node to fund their wallet.
	for i := range sts {
		err := waitForBlock(chainTip, sts[i])
		if err != nil {
			return err
		}

		// Mine a block. The next iteration of this loop will ensure that the
		// block propagates and does not get orphaned.
		block, err := sts[i].miner.AddBlock()
		if err != nil {
			return err
		}
		chainTip = block.ID()
	}

	// Wait until the chain tip has propagated to the first node.
	err = waitForBlock(chainTip, sts[0])
	if err != nil {
		return err
	}

	// Mine types.MaturityDelay more blocks from the final node to mine a
	// block, to guarantee that all nodes have had their payouts mature, such
	// that their wallets can begin spending immediately.
	for i := types.BlockHeight(0); i <= types.MaturityDelay; i++ {
		_, err := sts[0].miner.AddBlock()
		if err != nil {
			return err
		}
	}

	// Block until every node has the full chain.
	_, err = synchronizationCheck(sts)
	return err
}

// synchronizationCheck takes a bunch of server testers as input and checks
// that they all have the same current block as the first server tester. The
// first server tester needs to have the most recent block in order for the
// check to work.
func synchronizationCheck(sts []*serverTester) (types.BlockID, error) {
	// Prefer returning an error in the event of a zero-length server tester -
	// an error should be returned if the developer accidentally uses a nil
	// slice instead of whatever value was intended, and there's no reason to
	// check for synchronization if there aren't any nodes to be synchronized.
	if len(sts) == 0 {
		return types.BlockID{}, errors.New("no server testers provided")
	}

	// Wait until all nodes are on the same block
	for _, st := range sts[1:] {
		err := waitForBlock(sts[0].cs.CurrentBlock().ID(), st)
		if err != nil {
			return types.BlockID{}, err
		}
	}

	var cg ConsensusGET
	err := sts[0].getAPI("/consensus", &cg)
	if err != nil {
		return types.BlockID{}, err
	}
	leaderBlockID := cg.CurrentBlock
	for i := range sts {
		// Spin until the current block matches the leader block.
		success := false
		for j := 0; j < 100; j++ {
			err = sts[i].getAPI("/consensus", &cg)
			if err != nil {
				return types.BlockID{}, err
			}
			if cg.CurrentBlock == leaderBlockID {
				success = true
				break
			}
			time.Sleep(time.Millisecond * 100)
		}
		if !success {
			return types.BlockID{}, errors.New("synchronization check failed - nodes do not seem to be synchronized")
		}
	}
	return leaderBlockID, nil
}

// waitForBlock will block until the provided chain tip is the most recent
// block in the provided testing node.
func waitForBlock(chainTip types.BlockID, st *serverTester) error {
	var cg ConsensusGET
	success := false
	for j := 0; j < 100; j++ {
		err := st.getAPI("/consensus", &cg)
		if err != nil {
			return err
		}
		if cg.CurrentBlock == chainTip {
			success = true
			break
		}
		time.Sleep(time.Millisecond * 100)
	}
	if !success {
		return errors.New("node never reached the correct chain tip")
	}
	return nil
}
<commit_msg>use build.Retry to fix NDF<commit_after>package api

// ecosystem_helpers_test.go has a bunch of helper functions to make setting up
// large ecosystem tests easier.
//
// List of helper functions:
//    addStorageToAllHosts // adds a storage folder to every host
//    announceAllHosts     // announce all hosts to the network (and mine a block)
//    fullyConnectNodes    // connects each server tester to all the others
//    fundAllNodes         // mines blocks until all server testers have money
//    synchronizationCheck // checks that all server testers have the same recent block
//    waitForBlock         // block until the provided block is the most recent block for all server testers

import (
	"errors"
	"fmt"
	"net/url"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/types"
)

// addStorageToAllHosts adds a storage folder with a bunch of storage to each
// host.
func addStorageToAllHosts(sts []*serverTester) error {
	for _, st := range sts {
		values := url.Values{}
		values.Set("path", st.dir)
		values.Set("size", "1048576")
		err := st.stdPostAPI("/host/storage/folders/add", values)
		if err != nil {
			return err
		}
	}
	return nil
}

// announceAllHosts will announce every host in the tester set to the
// blockchain.
func announceAllHosts(sts []*serverTester) error {
	// Check that all announcements will be on the same chain.
	_, err := synchronizationCheck(sts)
	if err != nil {
		return err
	}

	// Grab the initial transaction pool size to know how many total transactions
	// there should be after announcement.
	initialTpoolSize := len(sts[0].tpool.TransactionList())

	// Announce each host.
	for _, st := range sts {
		// Set the host to be accepting contracts.
		acceptingContractsValues := url.Values{}
		acceptingContractsValues.Set("acceptingcontracts", "true")
		err = st.stdPostAPI("/host", acceptingContractsValues)
		if err != nil {
			return err
		}

		// Fetch the host net address.
		var hg HostGET
		err = st.getAPI("/host", &hg)
		if err != nil {
			return err
		}

		// Make the announcement.
		announceValues := url.Values{}
		announceValues.Set("address", string(hg.ExternalSettings.NetAddress))
		err = st.stdPostAPI("/host/announce", announceValues)
		if err != nil {
			return err
		}
	}

	// Wait until all of the transactions have propagated to all of the nodes.
	//
	// TODO: Replace this direct transaction pool call with a call to the
	// /transactionpool endpoint.
	//
	// TODO: At some point the number of transactions needed to make an
	// announcement may change. Currently its 2.
	for i := 0; i < 50; i++ {
		if len(sts[0].tpool.TransactionList()) == len(sts)*2+initialTpoolSize {
			break
		}
		time.Sleep(time.Millisecond * 100)
	}
	if len(sts[0].tpool.TransactionList()) < len(sts)*2+initialTpoolSize {
		return fmt.Errorf("Host announcements do not seem to have propagated to the leader's tpool: %v, %v, %v", len(sts), len(sts[0].tpool.TransactionList())+initialTpoolSize, initialTpoolSize)
	}

	// Mine a block and then wait for all of the nodes to syncrhonize to it.
	_, err = sts[0].miner.AddBlock()
	if err != nil {
		return err
	}
	// Block until the block propagated to all nodes
	for _, st := range sts[1:] {
		err = waitForBlock(sts[0].cs.CurrentBlock().ID(), st)
		if err != nil {
			return (err)
		}
	}
	// Check if all nodes are on the same block now
	_, err = synchronizationCheck(sts)
	if err != nil {
		return err
	}

	// Block until every node has completed the scan of every other node, so
	// that each node has a full hostdb.
	for _, st := range sts {
		err := build.Retry(600, 100*time.Millisecond, func() error {
			var ah HostdbActiveGET
			err = st.getAPI("/hostdb/active", &ah)
			if err != nil {
				return err
			}
			if len(ah.Hosts) < len(sts) {
				return errors.New("one of the nodes hostdbs was unable to find at least one host announcement")
			}
			return nil
		})
		if err != nil {
			return err
		}
	}
	return nil
}

// fullyConnectNodes takes a bunch of tester nodes and connects each to the
// other, creating a fully connected graph so that everyone is on the same
// chain.
//
// After connecting the nodes, it verifies that all the nodes have
// synchronized.
func fullyConnectNodes(sts []*serverTester) error {
	for i, sta := range sts {
		var gg GatewayGET
		err := sta.getAPI("/gateway", &gg)
		if err != nil {
			return err
		}

		// Connect this node to every other node.
		for _, stb := range sts[i+1:] {
			// Try connecting to the other node until both have the other in
			// their peer list.
			err = retry(100, time.Millisecond*100, func() error {
				// NOTE: this check depends on string-matching an error in the
				// gateway. If that error changes at all, this string will need to
				// be updated.
				err := stb.stdPostAPI("/gateway/connect/"+string(gg.NetAddress), nil)
				if err != nil && err.Error() != "already connected to this peer" {
					return err
				}

				// Check that the gateways are connected.
				bToA := false
				aToB := false
				var ggb GatewayGET
				err = stb.getAPI("/gateway", &ggb)
				if err != nil {
					return err
				}
				for _, peer := range ggb.Peers {
					if peer.NetAddress == gg.NetAddress {
						bToA = true
						break
					}
				}
				err = sta.getAPI("/gateway", &gg)
				if err != nil {
					return err
				}
				for _, peer := range gg.Peers {
					if peer.NetAddress == ggb.NetAddress {
						aToB = true
						break
					}
				}
				if !aToB || !bToA {
					return fmt.Errorf("called connect between two nodes, but they are not peers: %v %v %v %v %v %v", aToB, bToA, gg.NetAddress, ggb.NetAddress, gg.Peers, ggb.Peers)
				}
				return nil

			})
			if err != nil {
				return err
			}
		}
	}

	// Perform a synchronization check.
	_, err := synchronizationCheck(sts)
	return err
}

// fundAllNodes will make sure that each node has mined a block in the longest
// chain, then will mine enough blocks that the miner payouts manifest in the
// wallets of each node.
func fundAllNodes(sts []*serverTester) error {
	// Check that all of the nodes are synchronized.
	chainTip, err := synchronizationCheck(sts)
	if err != nil {
		return err
	}

	// Mine a block for each node to fund their wallet.
	for i := range sts {
		err := waitForBlock(chainTip, sts[i])
		if err != nil {
			return err
		}

		// Mine a block. The next iteration of this loop will ensure that the
		// block propagates and does not get orphaned.
		block, err := sts[i].miner.AddBlock()
		if err != nil {
			return err
		}
		chainTip = block.ID()
	}

	// Wait until the chain tip has propagated to the first node.
	err = waitForBlock(chainTip, sts[0])
	if err != nil {
		return err
	}

	// Mine types.MaturityDelay more blocks from the final node to mine a
	// block, to guarantee that all nodes have had their payouts mature, such
	// that their wallets can begin spending immediately.
	for i := types.BlockHeight(0); i <= types.MaturityDelay; i++ {
		_, err := sts[0].miner.AddBlock()
		if err != nil {
			return err
		}
	}

	// Block until every node has the full chain.
	_, err = synchronizationCheck(sts)
	return err
}

// synchronizationCheck takes a bunch of server testers as input and checks
// that they all have the same current block as the first server tester. The
// first server tester needs to have the most recent block in order for the
// check to work.
func synchronizationCheck(sts []*serverTester) (types.BlockID, error) {
	// Prefer returning an error in the event of a zero-length server tester -
	// an error should be returned if the developer accidentally uses a nil
	// slice instead of whatever value was intended, and there's no reason to
	// check for synchronization if there aren't any nodes to be synchronized.
	if len(sts) == 0 {
		return types.BlockID{}, errors.New("no server testers provided")
	}

	// Wait until all nodes are on the same block
	for _, st := range sts[1:] {
		err := waitForBlock(sts[0].cs.CurrentBlock().ID(), st)
		if err != nil {
			return types.BlockID{}, err
		}
	}

	var cg ConsensusGET
	err := sts[0].getAPI("/consensus", &cg)
	if err != nil {
		return types.BlockID{}, err
	}
	leaderBlockID := cg.CurrentBlock
	for i := range sts {
		// Spin until the current block matches the leader block.
		success := false
		for j := 0; j < 100; j++ {
			err = sts[i].getAPI("/consensus", &cg)
			if err != nil {
				return types.BlockID{}, err
			}
			if cg.CurrentBlock == leaderBlockID {
				success = true
				break
			}
			time.Sleep(time.Millisecond * 100)
		}
		if !success {
			return types.BlockID{}, errors.New("synchronization check failed - nodes do not seem to be synchronized")
		}
	}
	return leaderBlockID, nil
}

// waitForBlock will block until the provided chain tip is the most recent
// block in the provided testing node.
func waitForBlock(chainTip types.BlockID, st *serverTester) error {
	var cg ConsensusGET
	success := false
	for j := 0; j < 100; j++ {
		err := st.getAPI("/consensus", &cg)
		if err != nil {
			return err
		}
		if cg.CurrentBlock == chainTip {
			success = true
			break
		}
		time.Sleep(time.Millisecond * 100)
	}
	if !success {
		return errors.New("node never reached the correct chain tip")
	}
	return nil
}
<|endoftext|>


<commit_before>// Author  Raido Pahtma
// License MIT

// Package reportreceiver ReportReceiver
package reportreceiver

import (
	"errors"
	"fmt"
	"time"

	"github.com/proactivity-lab/go-loggers"
	"github.com/proactivity-lab/go-moteconnection"
)

const AMID_REPORTS = 9
const AM_DEFAULT_GROUP = 0x22

const HEADER_REPORTMESSAGE = 1
const HEADER_REPORTMESSAGE_ACK = 2

type ReportWriter interface {
	Append(*Report) error
}

type ReportMsg struct {
	Header   uint8
	Report   uint32
	Fragment uint8
	Total    uint8
	Data     []byte
}

type ReportMsgAck struct {
	Header  uint8
	Report  uint32
	Missing []uint8
}

type ReportData struct {
	Channel        uint8
	Id             uint32
	LocalTimeMilli uint32
	ClockTime      uint32
	Data           []byte
}

type Report struct {
	Source         moteconnection.AMAddr
	Report         uint32
	Channel        uint8
	Id             uint32
	LocalTimeMilli uint32 // MilliSeconds from boot
	ClockTime      uint32 // Seconds from the beginning of the century or 0xFFFFFFFF
	Data           []byte

	FirstRcvd time.Time
	LastRcvd  time.Time
	FragsRcvd int
}

type QueueElement struct {
	Source moteconnection.AMAddr
}

var (
	InputQueue = make(map[moteconnection.AMAddr]moteconnection.Message)
)

func (self *Report) StorageStringHeader() string {
	return "timestamp ff, timestamp lf, ADDR, reportnum, CHANNEL, reportid, clocktime, localtime, data"
}

func (self *Report) StorageString() string {
	tsl := "2006-01-02 15:04:05.999"
	return fmt.Sprintf("%s,%s,%s,%d,%02X,%d,%d,%d,%X", self.FirstRcvd.Format(tsl), self.LastRcvd.Format(tsl), self.Source, self.Report, self.Channel, self.Id, self.ClockTime, self.LocalTimeMilli, self.Data)
}

func (self *Report) String() string {
	return fmt.Sprintf("%s rprt %d([%02X]%d) %X", self.Source, self.Report, self.Channel, self.Id, self.Data)
}

type PartialReport struct {
	Source moteconnection.AMAddr
	Report uint32

	Total uint8

	FirstRcvd time.Time
	LastRcvd  time.Time
	FragsRcvd int

	LastContact time.Time

	Fragments map[uint8]ReportMsg
}

func (self *PartialReport) String() string {
	return fmt.Sprintf("%s rprt %d %d/%d", self.Source, self.Report, self.Total, len(self.Fragments))
}

func (self *PartialReport) AddFragment(fragment *ReportMsg) {
	self.Total = fragment.Total
	self.FragsRcvd++
	self.LastRcvd = time.Now().UTC()
	self.Fragments[fragment.Fragment] = *fragment
}

func (self *PartialReport) IsComplete() bool {
	if fragment, ok := self.Fragments[0]; ok {
		if len(self.Fragments) == int(fragment.Total) {
			return true
		}
	}
	return false
}

func (self *PartialReport) GetReport() (*Report, error) {
	if self.IsComplete() {
		report := new(Report)
		report.Source = self.Source
		report.Report = self.Report

		data := make([]byte, len(self.Fragments[0].Data)*len(self.Fragments))
		data = data[:0]
		for i := uint8(0); i < uint8(len(self.Fragments)); i++ {
			data = append(data, self.Fragments[i].Data...)
		}

		rpd := new(ReportData)
		if err := moteconnection.DeserializePacket(rpd, data); err != nil {
			return nil, err
		}

		report.Channel = rpd.Channel
		report.Id = rpd.Id
		report.ClockTime = rpd.ClockTime
		report.LocalTimeMilli = rpd.LocalTimeMilli
		report.Data = rpd.Data

		report.FirstRcvd = self.FirstRcvd
		report.LastRcvd = self.LastRcvd
		report.FragsRcvd = self.FragsRcvd

		return report, nil
	}
	return nil, errors.New("Report not complete")
}

func (self *PartialReport) Missing() []uint8 {
	missing := make([]uint8, self.Total)
	missing = missing[:0]
	for i := uint8(0); i < self.Total; i++ {
		if _, ok := self.Fragments[i]; !ok {
			missing = append(missing, i)
		}
	}
	return missing
}

func NewPartialReport(source moteconnection.AMAddr, rm *ReportMsg) *PartialReport {
	pr := new(PartialReport)
	pr.Source = source
	pr.Report = rm.Report
	pr.Fragments = make(map[uint8]ReportMsg)
	pr.FirstRcvd = time.Now().UTC()
	return pr
}

type ReportReceiver struct {
	loggers.DIWEloggers

	mconn        moteconnection.MoteConnection
	dsp          *moteconnection.MessageDispatcher
	receive      chan moteconnection.Packet
	reports      map[moteconnection.AMAddr]*PartialReport
	reportwriter ReportWriter
}

func NewReportReceiver(mconn moteconnection.MoteConnection, source moteconnection.AMAddr, group moteconnection.AMGroup) *ReportReceiver {
	rl := new(ReportReceiver)
	rl.InitLoggers()

	rl.receive = make(chan moteconnection.Packet)
	rl.reports = make(map[moteconnection.AMAddr]*PartialReport)

	rl.dsp = moteconnection.NewMessageDispatcher(moteconnection.NewMessage(group, source))
	rl.dsp.RegisterMessageReceiver(AMID_REPORTS, rl.receive)

	rl.mconn = mconn
	rl.mconn.AddDispatcher(rl.dsp)

	return rl
}

func (self *ReportReceiver) SendAck(destination moteconnection.AMAddr, report uint32, missing []uint8) {
	ack := new(ReportMsgAck)
	ack.Header = 2
	ack.Report = report
	ack.Missing = missing

	msg := self.dsp.NewMessage()
	msg.SetDestination(destination)
	msg.SetType(AMID_REPORTS)
	msg.Payload = moteconnection.SerializePacket(ack)

	self.mconn.Send(msg)
}

func (self *ReportReceiver) SetOutput(rw ReportWriter) {
	self.reportwriter = rw
}

func (self *ReportReceiver) Run() {
	self.Debug.Printf("run main loop\n")
	for {
		//self.Debug.Printf("Main Loop begin\n")
		select {
		case packet := <-self.receive:
			msg := packet.(*moteconnection.Message)
			self.Debug.Printf("%s\n", msg)
			if len(msg.Payload) > 0 {
				if msg.Payload[0] == HEADER_REPORTMESSAGE {
					rpm := new(ReportMsg)

					if err := moteconnection.DeserializePacket(rpm, msg.Payload); err != nil {
						self.Error.Printf("%s %s\n", msg, err)
						break
					}
					/*
						if _, ok := InputQueue[msg.Source()]; ok {
							self.Debug.Printf("Deleting from queue reset from: %s", msg.Source())
							delete(InputQueue, msg.Source())
						}
					*/
					if rpm.Report == 0 {
						self.Info.Printf("RESET %s\n", msg.Source())
						//InputQueue[msg.Source()] = *msg
						delete(self.reports, msg.Source())
					}

					pr, ok := self.reports[msg.Source()]
					if !ok {
						self.reports[msg.Source()] = NewPartialReport(msg.Source(), rpm)
						pr = self.reports[msg.Source()]
					}

					if rpm.Report != pr.Report {
						// Getting a different report than expected for some reason
						if rpm.Report > pr.Report { // Report skip
							if pr.IsComplete() == false {
								self.Debug.Printf("skip %s rprt %d->%d\n", msg.Source(), pr.Report, rpm.Report)
							} else {
								self.Debug.Printf("new %s rprt %d\n", msg.Source(), rpm.Report)
							}
							self.reports[msg.Source()] = NewPartialReport(msg.Source(), rpm)
							pr = self.reports[msg.Source()]
						} else { // Older report
							self.SendAck(msg.Source(), rpm.Report, nil)
							// TODO Should only ack like once really, if there are several fragments
							break
						}
					} else {
						if pr.IsComplete() {
							self.Debug.Printf("repeat %s rprt %d\n", msg.Source(), rpm.Report)
							self.SendAck(msg.Source(), rpm.Report, nil)
							break
						}
					}

					pr.AddFragment(rpm)
					self.Debug.Printf("%s\n", pr)

					if report, err := pr.GetReport(); err == nil {
						self.Info.Printf("%s\n", report)
						err := self.reportwriter.Append(report)
						if err != nil {
							self.Error.Printf("%s\n", err)
						}
						//sends ACK if all fragments have arrived
						self.SendAck(msg.Source(), rpm.Report, nil)
					} else {
						// self.SendAck(msg.Source(), rpm.Report, pr.Missing())
						// TODO Delay ack, still missing something
					}
				}
			}
		}
		//self.Debug.Println("Current Queue length:", len(InputQueue))
	}
}

// RunResetResender func
// Resends ACK packets to nodes which are awaiting for routing and have only sent RESET
func (self *ReportReceiver) RunResetResender() {
	self.Debug.Println("Run RESET ACK sender")
	for {
		time.Sleep(2 * time.Minute)
		for key, element := range self.reports {
			if 2 > element.Total {
				self.Debug.Printf("Reportlogger shall send ACK from reset queue")
				self.Debug.Printf("Sending to: %s", key)
				self.SendAck(key, element.Report, nil)
				break
			}
		}
	}
}

// RunMissingFragmentResender func
// Resends ACK packets to nodes which have not sent all fragments at current report
func (self *ReportReceiver) RunMissingFragmentResender() {
	self.Debug.Println("Run missing fragment sender")
	for {
		time.Sleep(1 * time.Minute)
		for key, element := range self.reports {
			if !element.IsComplete() {
				self.Debug.Printf("missing fragment, sending to: %s", key)
				self.SendAck(key, element.Report, nil)
				time.Sleep(5 * time.Second)
			}
		}

	}
}
<commit_msg>improved fragment receiving functionality, also removed unnecessary and commented code<commit_after>// Author  Raido Pahtma
// License MIT

// Package reportreceiver ReportReceiver
package reportreceiver

import (
	"errors"
	"fmt"
	"time"

	"github.com/proactivity-lab/go-loggers"
	"github.com/proactivity-lab/go-moteconnection"
)

const AMID_REPORTS = 9
const AM_DEFAULT_GROUP = 0x22

const HEADER_REPORTMESSAGE = 1
const HEADER_REPORTMESSAGE_ACK = 2

type ReportWriter interface {
	Append(*Report) error
}

type ReportMsg struct {
	Header   uint8
	Report   uint32
	Fragment uint8
	Total    uint8
	Data     []byte
}

type ReportMsgAck struct {
	Header  uint8
	Report  uint32
	Missing []uint8
}

type ReportData struct {
	Channel        uint8
	Id             uint32
	LocalTimeMilli uint32
	ClockTime      uint32
	Data           []byte
}

type Report struct {
	Source         moteconnection.AMAddr
	Report         uint32
	Channel        uint8
	Id             uint32
	LocalTimeMilli uint32 // MilliSeconds from boot
	ClockTime      uint32 // Seconds from the beginning of the century or 0xFFFFFFFF
	Data           []byte

	FirstRcvd time.Time
	LastRcvd  time.Time
	FragsRcvd int
}

type QueueElement struct {
	Source moteconnection.AMAddr
}

func (self *Report) StorageStringHeader() string {
	return "timestamp ff, timestamp lf, ADDR, reportnum, CHANNEL, reportid, clocktime, localtime, data"
}

func (self *Report) StorageString() string {
	tsl := "2006-01-02 15:04:05.999"
	return fmt.Sprintf("%s,%s,%s,%d,%02X,%d,%d,%d,%X", self.FirstRcvd.Format(tsl), self.LastRcvd.Format(tsl), self.Source, self.Report, self.Channel, self.Id, self.ClockTime, self.LocalTimeMilli, self.Data)
}

func (self *Report) String() string {
	return fmt.Sprintf("%s rprt %d([%02X]%d) %X", self.Source, self.Report, self.Channel, self.Id, self.Data)
}

type PartialReport struct {
	Source moteconnection.AMAddr
	Report uint32

	Total uint8

	FirstRcvd time.Time
	LastRcvd  time.Time
	FragsRcvd int

	LastContact time.Time

	Fragments map[uint8]ReportMsg
}

func (self *PartialReport) String() string {
	return fmt.Sprintf("%s rprt %d %d/%d", self.Source, self.Report, self.Total, len(self.Fragments))
}

func (self *PartialReport) AddFragment(fragment *ReportMsg) {
	self.Total = fragment.Total
	self.FragsRcvd++
	self.LastRcvd = time.Now().UTC()
	self.Fragments[fragment.Fragment] = *fragment
}

func (self *PartialReport) IsComplete() bool {
	if fragment, ok := self.Fragments[0]; ok {
		if len(self.Fragments) == int(fragment.Total) {
			return true
		}
	}
	return false
}

func (self *PartialReport) GetReport() (*Report, error) {
	if self.IsComplete() {
		report := new(Report)
		report.Source = self.Source
		report.Report = self.Report

		data := make([]byte, len(self.Fragments[0].Data)*len(self.Fragments))
		data = data[:0]
		for i := uint8(0); i < uint8(len(self.Fragments)); i++ {
			data = append(data, self.Fragments[i].Data...)
		}

		rpd := new(ReportData)
		if err := moteconnection.DeserializePacket(rpd, data); err != nil {
			return nil, err
		}

		report.Channel = rpd.Channel
		report.Id = rpd.Id
		report.ClockTime = rpd.ClockTime
		report.LocalTimeMilli = rpd.LocalTimeMilli
		report.Data = rpd.Data

		report.FirstRcvd = self.FirstRcvd
		report.LastRcvd = self.LastRcvd
		report.FragsRcvd = self.FragsRcvd

		return report, nil
	}
	return nil, errors.New("Report not complete")
}

func (self *PartialReport) Missing() []uint8 {
	missing := make([]uint8, self.Total)
	missing = missing[:0]
	for i := uint8(0); i < self.Total; i++ {
		if _, ok := self.Fragments[i]; !ok {
			missing = append(missing, i)
		}
	}
	return missing
}

func NewPartialReport(source moteconnection.AMAddr, rm *ReportMsg) *PartialReport {
	pr := new(PartialReport)
	pr.Source = source
	pr.Report = rm.Report
	pr.Fragments = make(map[uint8]ReportMsg)
	pr.FirstRcvd = time.Now().UTC()
	return pr
}

type ReportReceiver struct {
	loggers.DIWEloggers

	mconn        moteconnection.MoteConnection
	dsp          *moteconnection.MessageDispatcher
	receive      chan moteconnection.Packet
	reports      map[moteconnection.AMAddr]*PartialReport
	reportwriter ReportWriter
}

func NewReportReceiver(mconn moteconnection.MoteConnection, source moteconnection.AMAddr, group moteconnection.AMGroup) *ReportReceiver {
	rl := new(ReportReceiver)
	rl.InitLoggers()

	rl.receive = make(chan moteconnection.Packet)
	rl.reports = make(map[moteconnection.AMAddr]*PartialReport)

	rl.dsp = moteconnection.NewMessageDispatcher(moteconnection.NewMessage(group, source))
	rl.dsp.RegisterMessageReceiver(AMID_REPORTS, rl.receive)

	rl.mconn = mconn
	rl.mconn.AddDispatcher(rl.dsp)

	return rl
}

func (self *ReportReceiver) SendAck(destination moteconnection.AMAddr, report uint32, missing []uint8) {
	ack := new(ReportMsgAck)
	ack.Header = 2
	ack.Report = report
	ack.Missing = missing

	msg := self.dsp.NewMessage()
	msg.SetDestination(destination)
	msg.SetType(AMID_REPORTS)
	msg.Payload = moteconnection.SerializePacket(ack)

	self.mconn.Send(msg)
}

func (self *ReportReceiver) SetOutput(rw ReportWriter) {
	self.reportwriter = rw
}

func (self *ReportReceiver) Run() {
	self.Debug.Printf("run main loop\n")
	for {
		select {
		case packet := <-self.receive:
			msg := packet.(*moteconnection.Message)
			self.Debug.Printf("%s\n", msg)
			if len(msg.Payload) > 0 {
				if msg.Payload[0] == HEADER_REPORTMESSAGE {
					rpm := new(ReportMsg)

					if err := moteconnection.DeserializePacket(rpm, msg.Payload); err != nil {
						self.Error.Printf("%s %s\n", msg, err)
						break
					}
					if rpm.Report == 0 {
						self.Info.Printf("RESET %s\n", msg.Source())
						delete(self.reports, msg.Source())
					}

					pr, ok := self.reports[msg.Source()]
					if !ok {
						self.reports[msg.Source()] = NewPartialReport(msg.Source(), rpm)
						pr = self.reports[msg.Source()]
					}

					if rpm.Report != pr.Report {
						// Getting a different report than expected for some reason
						if rpm.Report > pr.Report { // Report skip
							if pr.IsComplete() == false {
								self.Debug.Printf("skip %s rprt %d->%d\n", msg.Source(), pr.Report, rpm.Report)
							} else {
								self.Debug.Printf("new %s rprt %d\n", msg.Source(), rpm.Report)
							}
							self.reports[msg.Source()] = NewPartialReport(msg.Source(), rpm)
							pr = self.reports[msg.Source()]
						} else { // Older report
							self.SendAck(msg.Source(), rpm.Report, nil)
							// TODO Should only ack like once really, if there are several fragments
							break
						}
					} else {
						if pr.IsComplete() {
							self.Debug.Printf("repeat %s rprt %d\n", msg.Source(), rpm.Report)
							self.SendAck(msg.Source(), rpm.Report, nil)
							break
						}
					}

					pr.AddFragment(rpm)
					self.Debug.Printf("%s\n", pr)
					if pr.IsComplete() {
						if report, err := pr.GetReport(); err == nil {
							self.Info.Printf("%s\n", report)
							err := self.reportwriter.Append(report)
							if err != nil {
								self.Error.Printf("%s\n", err)
							}
							//sends ACK if all fragments have arrived
							self.SendAck(msg.Source(), rpm.Report, nil)
						} else {
							self.Error.Printf("%s\n", err)
						}
					} else {
						self.Debug.Printf("Missing fragments!")
					}
				}
			}
		}
	}
}

// RunResetResender func
// Resends ACK packets to nodes which are awaiting for routing and have only sent RESET
func (self *ReportReceiver) RunResetResender() {
	self.Debug.Println("Run RESET ACK sender")
	for {
		time.Sleep(2 * time.Minute)
		for key, element := range self.reports {
			if 2 > element.Total {
				self.Debug.Printf("Reportlogger shall send ACK from reset queue")
				self.Debug.Printf("Sending to: %s", key)
				self.SendAck(key, element.Report, nil)
				break
			}
		}
	}
}

// RunMissingFragmentResender func
// Resends ACK packets to nodes which have not sent all fragments at current report
func (self *ReportReceiver) RunMissingFragmentResender() {
	self.Debug.Println("Run missing fragment sender")
	for {
		time.Sleep(1 * time.Minute)
		for key, element := range self.reports {
			if !element.IsComplete() {
				self.Debug.Printf("missing fragment, sending to: %s", key)
				self.SendAck(key, element.Report, nil)
				time.Sleep(5 * time.Second)
			}
		}

	}
}
<|endoftext|>


<commit_before>//  Copyright (c) 2017 Couchbase, Inc.
//
//  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 moss

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
	"testing"
	"time"
)

func TestSegmentKeysIndex(t *testing.T) {
	// Consider a moss segment with this data:
	// Buf: |key1|val1|key10|val10|key100|val100|key1000|val1000
	//      |key250|val250|key4000|val4000|key500|val500
	// Kvs: |4|4|0    |5|5|8      |6|6|18       |7|7|30
	//      |7|7|44       |7|7|56         |6|6|70

	keys := []string{
		"key1",
		"key10",
		"key100",
		"key1000",
		"key250",
		"key4000",
		"key500",
	}

	s := newSegmentKeysIndex(
		30, // quota
		7,  // number of keys
		6,  // average key size
	)

	for i, k := range keys {
		ret := s.add(i, []byte(k))
		if !ret {
			t.Errorf("Unexpected Add failure!")
		}
	}

	ret := s.add(0, []byte("key"))
	if ret {
		t.Errorf("Space shouldn't have been available!")
	}

	if s.numIndexableKeys != 3 {
		t.Errorf("Unexpected number of keys (%v) indexed!", s.numIndexableKeys)
	}

	data := []byte("key1key1000key500")

	if !bytes.Contains(s.data, data) {
		t.Errorf("Unexpected data in index array: %v", string(s.data))
	}

	if s.numKeys != 3 ||
		s.offsets[0] != 0 || s.offsets[1] != 4 || s.offsets[2] != 11 {
		t.Errorf("Unexpected content in offsets array!")
	}

	if s.hop != 3 {
		t.Errorf("Unexpected hop: %v!", s.hop)
	}

	var left, right int

	left, right = s.lookup([]byte("key1000"))
	if left != 3 || right != 4 {
		t.Errorf("Unexpected results for key1000")
	}

	left, right = s.lookup([]byte("key1"))
	if left != 0 || right != 1 {
		t.Errorf("Unexpected results for key1")
	}

	left, right = s.lookup([]byte("key500"))
	if left != 6 || right != 7 {
		t.Errorf("Unexpected results for key500")
	}

	left, right = s.lookup([]byte("key400"))
	if left != 3 || right != 6 {
		t.Errorf("Unexpected results for key4000")
	}

	left, right = s.lookup([]byte("key100"))
	if left != 0 || right != 3 {
		t.Errorf("Unexpected results for key100")
	}

	left, right = s.lookup([]byte("key0"))
	if left != 0 || right != 0 {
		t.Errorf("Unexpected results for key0")
	}

	left, right = s.lookup([]byte("key6"))
	if left != 6 || right != 7 {
		t.Errorf("Unexpected results for key6")
	}
}

type byNS []time.Duration

func (x byNS) Len() int           { return len(x) }
func (x byNS) Swap(i, j int)      { x[i], x[j] = x[j], x[i] }
func (x byNS) Less(i, j int) bool { return x[i].Nanoseconds() < x[j].Nanoseconds() }

type testResults struct {
	name       string
	batchsize  int
	writetime  time.Duration
	readtime   time.Duration
	fetchtimes []time.Duration
	mean       time.Duration
}

func TestSegmentKindBasicWithAndWithoutIndex(t *testing.T) {
	numItems := 1000000

	ch := make(chan *testResults)

	runTest := func(name string, batchSize, quota, minSize int) {
		tmpDir, _ := ioutil.TempDir("", "mossStore")
		defer os.RemoveAll(tmpDir)

		start := time.Now()

		so := DefaultStoreOptions
		so.CompactionSync = true
		so.SegmentKeysIndexMaxBytes = quota
		so.SegmentKeysIndexMinKeyBytes = minSize
		spo := StorePersistOptions{CompactionConcern: CompactionAllow}

		store, coll, er := OpenStoreCollection(tmpDir, so, spo)
		if er != nil || store == nil || coll == nil {
			t.Fatalf("error opening store collection: %v", tmpDir)
		}

		x := 0
		for i := 0; i < numItems; i = i + batchSize {
			ba, err := coll.NewBatch(batchSize, batchSize*512)
			if err != nil {
				t.Fatalf("error creating new batch: %v", err)
			}

			for j := i + batchSize - 1; j >= i; j-- {
				k := fmt.Sprintf("%08d", x)
				v := fmt.Sprintf("%128d", x)
				ba.Set([]byte(k), []byte(v))
				x++
			}
			err = coll.ExecuteBatch(ba, WriteOptions{})
			if err != nil {
				t.Fatalf("error executing batch: %v", err)
			}

			err = ba.Close()
			if err != nil {
				t.Fatalf("error closing batch: %v", err)
			}
		}

		writetime := time.Since(start)

		fetchtimes := make([]time.Duration, numItems/10)
		var aggregate int64

		start = time.Now()

		for i := 0; i < numItems/10; i++ {
			gstart := time.Now()
			val, err := coll.Get([]byte(fmt.Sprintf("%08d", i*10)), ReadOptions{})
			fetchtimes[i] = time.Since(gstart)
			expect := fmt.Sprintf("%128d", i*10)
			if err != nil || string(val) != expect {
				t.Fatalf("Unexpected error: %v / Vals mismatch: '%v' != '%v'",
					err, string(val), expect)
			}
			aggregate += fetchtimes[i].Nanoseconds()
		}

		readtime := time.Since(start)

		sort.Sort(byNS(fetchtimes))

		mean := time.Duration(aggregate/int64(len(fetchtimes))) * time.Nanosecond

		ch <- &testResults{
			name:       name,
			batchsize:  batchSize,
			writetime:  writetime,
			readtime:   readtime,
			fetchtimes: fetchtimes,
			mean:       mean,
		}
	}

	go runTest("WITHOUT SEG INDEX", 100, 0, 0)
	go runTest("WITH SEG INDEX", 100, 100000, 1000000)
	go runTest("WITHOUT SEG INDEX", 10000, 0, 0)
	go runTest("WITH SEG INDEX", 10000, 100000, 1000000)

	var results []*testResults
	for i := 0; i < 4; i++ {
		results = append(results, <-ch)
	}

	fmt.Printf("%17v (numItems: %10v) %25v %v\n",
		" ", numItems,
		" ", "<-------------------- Fetch times -------------------->")
	fmt.Printf("%17v  %9v  %15v  %15v  %10v  %10v  %10v  %10v  %10v\n",
		" ", "BatchSize", "Writetime", "Readtime",
		"Mean", "Median", "90th", "95th", "99th")
	for _, result := range results {
		fmt.Printf("%17v  %9v  %15v  %15v  %10v  %10v  %10v  %10v  %10v\n",
			result.name,
			result.batchsize,
			result.writetime,
			result.readtime,
			result.mean,
			result.fetchtimes[len(result.fetchtimes)/2],
			result.fetchtimes[90*len(result.fetchtimes)/100],
			result.fetchtimes[95*len(result.fetchtimes)/100],
			result.fetchtimes[99*len(result.fetchtimes)/100])
	}
}
<commit_msg>MB-25705: Rounding off read/write times to 3 decimal spaces<commit_after>//  Copyright (c) 2017 Couchbase, Inc.
//
//  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 moss

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
	"testing"
	"time"
)

func TestSegmentKeysIndex(t *testing.T) {
	// Consider a moss segment with this data:
	// Buf: |key1|val1|key10|val10|key100|val100|key1000|val1000
	//      |key250|val250|key4000|val4000|key500|val500
	// Kvs: |4|4|0    |5|5|8      |6|6|18       |7|7|30
	//      |7|7|44       |7|7|56         |6|6|70

	keys := []string{
		"key1",
		"key10",
		"key100",
		"key1000",
		"key250",
		"key4000",
		"key500",
	}

	s := newSegmentKeysIndex(
		30, // quota
		7,  // number of keys
		6,  // average key size
	)

	for i, k := range keys {
		ret := s.add(i, []byte(k))
		if !ret {
			t.Errorf("Unexpected Add failure!")
		}
	}

	ret := s.add(0, []byte("key"))
	if ret {
		t.Errorf("Space shouldn't have been available!")
	}

	if s.numIndexableKeys != 3 {
		t.Errorf("Unexpected number of keys (%v) indexed!", s.numIndexableKeys)
	}

	data := []byte("key1key1000key500")

	if !bytes.Contains(s.data, data) {
		t.Errorf("Unexpected data in index array: %v", string(s.data))
	}

	if s.numKeys != 3 ||
		s.offsets[0] != 0 || s.offsets[1] != 4 || s.offsets[2] != 11 {
		t.Errorf("Unexpected content in offsets array!")
	}

	if s.hop != 3 {
		t.Errorf("Unexpected hop: %v!", s.hop)
	}

	var left, right int

	left, right = s.lookup([]byte("key1000"))
	if left != 3 || right != 4 {
		t.Errorf("Unexpected results for key1000")
	}

	left, right = s.lookup([]byte("key1"))
	if left != 0 || right != 1 {
		t.Errorf("Unexpected results for key1")
	}

	left, right = s.lookup([]byte("key500"))
	if left != 6 || right != 7 {
		t.Errorf("Unexpected results for key500")
	}

	left, right = s.lookup([]byte("key400"))
	if left != 3 || right != 6 {
		t.Errorf("Unexpected results for key4000")
	}

	left, right = s.lookup([]byte("key100"))
	if left != 0 || right != 3 {
		t.Errorf("Unexpected results for key100")
	}

	left, right = s.lookup([]byte("key0"))
	if left != 0 || right != 0 {
		t.Errorf("Unexpected results for key0")
	}

	left, right = s.lookup([]byte("key6"))
	if left != 6 || right != 7 {
		t.Errorf("Unexpected results for key6")
	}
}

type byNS []time.Duration

func (x byNS) Len() int           { return len(x) }
func (x byNS) Swap(i, j int)      { x[i], x[j] = x[j], x[i] }
func (x byNS) Less(i, j int) bool { return x[i].Nanoseconds() < x[j].Nanoseconds() }

type testResults struct {
	name       string
	batchsize  int
	writetime  time.Duration
	readtime   time.Duration
	fetchtimes []time.Duration
	mean       time.Duration
}

func TestSegmentKindBasicWithAndWithoutIndex(t *testing.T) {
	numItems := 1000000

	ch := make(chan *testResults)

	runTest := func(name string, batchSize, quota, minSize int) {
		tmpDir, _ := ioutil.TempDir("", "mossStore")
		defer os.RemoveAll(tmpDir)

		start := time.Now()

		so := DefaultStoreOptions
		so.CompactionSync = true
		so.SegmentKeysIndexMaxBytes = quota
		so.SegmentKeysIndexMinKeyBytes = minSize
		spo := StorePersistOptions{CompactionConcern: CompactionAllow}

		store, coll, er := OpenStoreCollection(tmpDir, so, spo)
		if er != nil || store == nil || coll == nil {
			t.Fatalf("error opening store collection: %v", tmpDir)
		}

		x := 0
		for i := 0; i < numItems; i = i + batchSize {
			ba, err := coll.NewBatch(batchSize, batchSize*512)
			if err != nil {
				t.Fatalf("error creating new batch: %v", err)
			}

			for j := i + batchSize - 1; j >= i; j-- {
				k := fmt.Sprintf("%08d", x)
				v := fmt.Sprintf("%128d", x)
				ba.Set([]byte(k), []byte(v))
				x++
			}
			err = coll.ExecuteBatch(ba, WriteOptions{})
			if err != nil {
				t.Fatalf("error executing batch: %v", err)
			}

			err = ba.Close()
			if err != nil {
				t.Fatalf("error closing batch: %v", err)
			}
		}

		writetime := time.Since(start)

		fetchtimes := make([]time.Duration, numItems/10)
		var aggregate int64

		start = time.Now()

		for i := 0; i < numItems/10; i++ {
			gstart := time.Now()
			val, err := coll.Get([]byte(fmt.Sprintf("%08d", i*10)), ReadOptions{})
			fetchtimes[i] = time.Since(gstart)
			expect := fmt.Sprintf("%128d", i*10)
			if err != nil || string(val) != expect {
				t.Fatalf("Unexpected error: %v / Vals mismatch: '%v' != '%v'",
					err, string(val), expect)
			}
			aggregate += fetchtimes[i].Nanoseconds()
		}

		readtime := time.Since(start)

		sort.Sort(byNS(fetchtimes))

		mean := time.Duration(aggregate/int64(len(fetchtimes))) * time.Nanosecond

		ch <- &testResults{
			name:       name,
			batchsize:  batchSize,
			writetime:  writetime,
			readtime:   readtime,
			fetchtimes: fetchtimes,
			mean:       mean,
		}
	}

	go runTest("WITHOUT SEG INDEX", 100, 0, 0)
	go runTest("WITH SEG INDEX", 100, 100000, 1000000)
	go runTest("WITHOUT SEG INDEX", 10000, 0, 0)
	go runTest("WITH SEG INDEX", 10000, 100000, 1000000)

	var results []*testResults
	for i := 0; i < 4; i++ {
		results = append(results, <-ch)
	}

	fmt.Printf("%17v (numItems: %10v) %24v %v\n",
		" ", numItems,
		" ", "<-------------------- Fetch times -------------------->")
	fmt.Printf("%17v  %9v  %15v  %14v  %10v  %10v  %10v  %10v  %10v\n",
		" ", "BatchSize", "Writetime(s)", "Readtime(s)",
		"Mean", "Median", "90th", "95th", "99th")
	for _, result := range results {
		fmt.Printf("%17v  %9v  %15.3v  %14.3v  %10v  %10v  %10v  %10v  %10v\n",
			result.name,
			result.batchsize,
			result.writetime.Seconds(),
			result.readtime.Seconds(),
			result.mean,
			result.fetchtimes[len(result.fetchtimes)/2],
			result.fetchtimes[90*len(result.fetchtimes)/100],
			result.fetchtimes[95*len(result.fetchtimes)/100],
			result.fetchtimes[99*len(result.fetchtimes)/100])
	}
}
<|endoftext|>