<commit_before>package main

import (
	"crypto/sha256"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sns"
	"go.mozilla.org/autograph/formats"
	"go.mozilla.org/autograph/signer/apk"
	"go.mozilla.org/autograph/signer/apk2"
	"go.mozilla.org/autograph/signer/contentsignature"
	"go.mozilla.org/autograph/signer/contentsignaturepki"
	"go.mozilla.org/autograph/signer/genericrsa"
	"go.mozilla.org/autograph/signer/gpg2"
	"go.mozilla.org/autograph/signer/mar"
	"go.mozilla.org/autograph/signer/pgp"
	"go.mozilla.org/autograph/signer/rsapss"
	"go.mozilla.org/autograph/signer/xpi"
	"go.mozilla.org/hawk"
)

type configuration struct {
	url           string
	monitoringKey string
	env           string
	rootHash      string
	truststore    *x509.CertPool
	// hash keystore for verifying XPI dep signers
	depRootHash   string
	depTruststore *x509.CertPool
}

var conf configuration

const inputdata string = "AUTOGRAPH MONITORING"

var softNotifCache map[string]time.Time

func init() {
	// create a cache to avoid sending the same notifications over and over
	softNotifCache = make(map[string]time.Time)
}

func main() {
	conf.url = os.Getenv("AUTOGRAPH_URL")
	if conf.url == "" {
		log.Fatal("AUTOGRAPH_URL must be set to the base url of the autograph service")
	}
	conf.monitoringKey = os.Getenv("AUTOGRAPH_KEY")
	if conf.monitoringKey == "" {
		log.Fatal("AUTOGRAPH_KEY must be set to the api monitoring key")
	}
	conf.env = os.Getenv("AUTOGRAPH_ENV")
	switch conf.env {
	case "stage":
		conf.rootHash = firefoxPkiStageRootHash
		conf.truststore = x509.NewCertPool()
		conf.truststore.AppendCertsFromPEM([]byte(firefoxPkiStageRoot))
		conf.depRootHash = ""
		conf.depTruststore = nil
	case "prod":
		conf.rootHash = firefoxPkiProdRootHash
		conf.truststore = x509.NewCertPool()
		conf.truststore.AppendCertsFromPEM([]byte(firefoxPkiProdRoot))
		conf.depRootHash = firefoxPkiStageRootHash
		conf.depTruststore = x509.NewCertPool()
		conf.depTruststore.AppendCertsFromPEM([]byte(firefoxPkiStageRoot))
	default:
		conf.rootHash = "5E36F214DE823F8B299689235F0341ACAFA075AF82CB4CD4307C3DB343392AFE"
		conf.truststore = nil
		conf.depRootHash = ""
		conf.depTruststore = nil
	}
	if os.Getenv("AUTOGRAPH_ROOT_HASH") != "" {
		conf.rootHash = os.Getenv("AUTOGRAPH_ROOT_HASH")
	}
	if os.Getenv("LAMBDA_TASK_ROOT") != "" {
		// we are inside a lambda environment so run as lambda
		lambda.Start(Handler)
	} else {
		err := Handler()
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
			os.Exit(1)
		}
		os.Exit(0)
	}
}

// Handler contacts the autograph service and verifies all monitoring signatures
func Handler() (err error) {
	log.Println("Retrieving monitoring data from", conf.url)
	req, err := http.NewRequest("GET", conf.url+"__monitor__", nil)
	if err != nil {
		return
	}

	// For client requests, setting this field prevents re-use of
	// TCP connections between requests to the same hosts, as if
	// Transport.DisableKeepAlives were set.
	req.Close = true

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", makeAuthHeader(req, "monitor", conf.monitoringKey))
	cli := &http.Client{}
	resp, err := cli.Do(req)
	if err != nil || resp == nil {
		return
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}

	if resp.StatusCode != http.StatusCreated {
		return fmt.Errorf("Request failed with %s: %s", resp.Status, body)
	}

	// verify that we got a proper signature response, with valid signatures
	var responses []formats.SignatureResponse
	err = json.Unmarshal(body, &responses)
	if err != nil {
		return
	}
	failed := false
	var failures []error
	for i, response := range responses {
		switch response.Type {
		case contentsignature.Type:
			log.Printf("Verifying content signature from signer %q", response.SignerID)
			err = verifyContentSignature(response)
		case contentsignaturepki.Type:
			log.Printf("Verifying content signature pki from signer %q", response.SignerID)
			err = contentsignaturepki.Verify(response.X5U, response.Signature, []byte(inputdata))
		case xpi.Type:
			log.Printf("Verifying XPI signature from signer %q", response.SignerID)
			err = verifyXPISignature(response.Signature)
		case apk.Type:
			log.Printf("Verifying APK signature from signer %q", response.SignerID)
			err = verifyAPKSignature(response.Signature)
		case apk2.Type:
			// we don't verify apk2 signatures because they can only be obtained on valid
			// APK files, which is too big to fit in the monitoring logic
			log.Printf("Skipping verification of APK2 signature from signer %q (we can't verify those)", response.SignerID)
			continue
		case mar.Type:
			log.Printf("Verifying MAR signature from signer %q", response.SignerID)
			err = verifyMARSignature(response.Signature, response.PublicKey)
		case genericrsa.Type:
			log.Printf("Verifying RSA signature from signer %q", response.SignerID)
			err = genericrsa.VerifyGenericRsaSignatureResponse([]byte(inputdata), response)
		case rsapss.Type:
			log.Printf("Verifying RSA-PSS signature from signer %q", response.SignerID)
			err = verifyRsapssSignature(response.Signature, response.PublicKey)
		case pgp.Type, gpg2.Type:
			// we don't verify pgp signatures because that requires building a keyring
			// using the public key which is hard to do using the current openpgp package
			log.Printf("Skipping verification of PGP signature from signer %q (we can't verify those)", response.SignerID)
			continue
		default:
			err = fmt.Errorf("unknown signature type %q", response.Type)
		}
		if err != nil {
			failed = true
			log.Printf("Response %d from signer %q does not pass: %v", i, response.SignerID, err)
			log.Printf("Response was: %+v", response)
			failures = append(failures, err)
		} else {
			log.Printf("Response %d from signer %q passes verification", i, response.SignerID)
		}
	}
	if failed {
		failure := "Errors found during monitoring:"
		for i, fail := range failures {
			failure += fmt.Sprintf("\n%d. %s", i+1, fail.Error())
		}
		return fmt.Errorf(failure)
	}
	log.Println("All signature responses passed, monitoring OK")
	return
}

func makeAuthHeader(req *http.Request, user, token string) string {
	auth := hawk.NewRequestAuth(req,
		&hawk.Credentials{
			ID:   user,
			Key:  token,
			Hash: sha256.New},
		0)
	auth.Ext = fmt.Sprintf("%d", time.Now().Nanosecond())
	payloadhash := auth.PayloadHash("application/json")
	payloadhash.Write([]byte(""))
	auth.SetHash(payloadhash)
	return auth.RequestHeader()
}

// send a message to a predefined sns topic
func sendSoftNotification(id string, format string, a ...interface{}) error {
	if ts, ok := softNotifCache[id]; ok {
		// don't send dup notifications for 24 hours
		if ts.Add(24 * time.Hour).After(time.Now()) {
			log.Printf("silencing soft notification ID %s", id)
			return nil
		}
	}
	if os.Getenv("LAMBDA_TASK_ROOT") == "" || os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS") == "" {
		// We're not running in lambda or the conf isnt ready so don't try to publish to SQS
		log.Printf("soft notification ID %s: %s", id, fmt.Sprintf(format, a...))
		return nil
	}

	svc := sns.New(session.New())
	params := &sns.PublishInput{
		Message:  aws.String(fmt.Sprintf(format, a...)),
		TopicArn: aws.String(os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS")),
	}
	_, err := svc.Publish(params)
	if err != nil {
		return err
	}
	log.Printf("Soft notification send to %q with body: %s", os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS"), *params.Message)
	// add the notification to the cache
	softNotifCache[id] = time.Now()
	return nil
}
<commit_msg>in autograph monitor, force garbage collection before exiting<commit_after>package main

import (
	"crypto/sha256"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"runtime"
	"time"

	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sns"
	"go.mozilla.org/autograph/formats"
	"go.mozilla.org/autograph/signer/apk"
	"go.mozilla.org/autograph/signer/apk2"
	"go.mozilla.org/autograph/signer/contentsignature"
	"go.mozilla.org/autograph/signer/contentsignaturepki"
	"go.mozilla.org/autograph/signer/genericrsa"
	"go.mozilla.org/autograph/signer/gpg2"
	"go.mozilla.org/autograph/signer/mar"
	"go.mozilla.org/autograph/signer/pgp"
	"go.mozilla.org/autograph/signer/rsapss"
	"go.mozilla.org/autograph/signer/xpi"
	"go.mozilla.org/hawk"
)

type configuration struct {
	url           string
	monitoringKey string
	env           string
	rootHash      string
	truststore    *x509.CertPool
	// hash keystore for verifying XPI dep signers
	depRootHash   string
	depTruststore *x509.CertPool
}

var conf configuration

const inputdata string = "AUTOGRAPH MONITORING"

var softNotifCache map[string]time.Time

func init() {
	// create a cache to avoid sending the same notifications over and over
	softNotifCache = make(map[string]time.Time)
}

func main() {
	conf.url = os.Getenv("AUTOGRAPH_URL")
	if conf.url == "" {
		log.Fatal("AUTOGRAPH_URL must be set to the base url of the autograph service")
	}
	conf.monitoringKey = os.Getenv("AUTOGRAPH_KEY")
	if conf.monitoringKey == "" {
		log.Fatal("AUTOGRAPH_KEY must be set to the api monitoring key")
	}
	conf.env = os.Getenv("AUTOGRAPH_ENV")
	switch conf.env {
	case "stage":
		conf.rootHash = firefoxPkiStageRootHash
		conf.truststore = x509.NewCertPool()
		conf.truststore.AppendCertsFromPEM([]byte(firefoxPkiStageRoot))
		conf.depRootHash = ""
		conf.depTruststore = nil
	case "prod":
		conf.rootHash = firefoxPkiProdRootHash
		conf.truststore = x509.NewCertPool()
		conf.truststore.AppendCertsFromPEM([]byte(firefoxPkiProdRoot))
		conf.depRootHash = firefoxPkiStageRootHash
		conf.depTruststore = x509.NewCertPool()
		conf.depTruststore.AppendCertsFromPEM([]byte(firefoxPkiStageRoot))
	default:
		conf.rootHash = "5E36F214DE823F8B299689235F0341ACAFA075AF82CB4CD4307C3DB343392AFE"
		conf.truststore = nil
		conf.depRootHash = ""
		conf.depTruststore = nil
	}
	if os.Getenv("AUTOGRAPH_ROOT_HASH") != "" {
		conf.rootHash = os.Getenv("AUTOGRAPH_ROOT_HASH")
	}
	if os.Getenv("LAMBDA_TASK_ROOT") != "" {
		// we are inside a lambda environment so run as lambda
		lambda.Start(Handler)
	} else {
		err := Handler()
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
			os.Exit(1)
		}
		os.Exit(0)
	}
}

// Handler is a wrapper around monitor() that performs garbage collection
// before returning
func Handler() (err error) {
	defer func() {
		// force gc run
		// https://bugzilla.mozilla.org/show_bug.cgi?id=1621133
		t1 := time.Now()
		runtime.GC()
		log.Println("Garbage collected in", time.Now().Sub(t1))
	}()
	return monitor()
}

// monitor contacts the autograph service and verifies all monitoring signatures
func monitor() (err error) {
	log.Println("Retrieving monitoring data from", conf.url)
	req, err := http.NewRequest("GET", conf.url+"__monitor__", nil)
	if err != nil {
		return
	}

	// For client requests, setting this field prevents re-use of
	// TCP connections between requests to the same hosts, as if
	// Transport.DisableKeepAlives were set.
	req.Close = true

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", makeAuthHeader(req, "monitor", conf.monitoringKey))
	cli := &http.Client{}
	resp, err := cli.Do(req)
	if err != nil || resp == nil {
		return
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}

	if resp.StatusCode != http.StatusCreated {
		return fmt.Errorf("Request failed with %s: %s", resp.Status, body)
	}

	// verify that we got a proper signature response, with valid signatures
	var responses []formats.SignatureResponse
	err = json.Unmarshal(body, &responses)
	if err != nil {
		return
	}
	failed := false
	var failures []error
	for i, response := range responses {
		switch response.Type {
		case contentsignature.Type:
			log.Printf("Verifying content signature from signer %q", response.SignerID)
			err = verifyContentSignature(response)
		case contentsignaturepki.Type:
			log.Printf("Verifying content signature pki from signer %q", response.SignerID)
			err = contentsignaturepki.Verify(response.X5U, response.Signature, []byte(inputdata))
		case xpi.Type:
			log.Printf("Verifying XPI signature from signer %q", response.SignerID)
			err = verifyXPISignature(response.Signature)
		case apk.Type:
			log.Printf("Verifying APK signature from signer %q", response.SignerID)
			err = verifyAPKSignature(response.Signature)
		case apk2.Type:
			// we don't verify apk2 signatures because they can only be obtained on valid
			// APK files, which is too big to fit in the monitoring logic
			log.Printf("Skipping verification of APK2 signature from signer %q (we can't verify those)", response.SignerID)
			continue
		case mar.Type:
			log.Printf("Verifying MAR signature from signer %q", response.SignerID)
			err = verifyMARSignature(response.Signature, response.PublicKey)
		case genericrsa.Type:
			log.Printf("Verifying RSA signature from signer %q", response.SignerID)
			err = genericrsa.VerifyGenericRsaSignatureResponse([]byte(inputdata), response)
		case rsapss.Type:
			log.Printf("Verifying RSA-PSS signature from signer %q", response.SignerID)
			err = verifyRsapssSignature(response.Signature, response.PublicKey)
		case pgp.Type, gpg2.Type:
			// we don't verify pgp signatures because that requires building a keyring
			// using the public key which is hard to do using the current openpgp package
			log.Printf("Skipping verification of PGP signature from signer %q (we can't verify those)", response.SignerID)
			continue
		default:
			err = fmt.Errorf("unknown signature type %q", response.Type)
		}
		if err != nil {
			failed = true
			log.Printf("Response %d from signer %q does not pass: %v", i, response.SignerID, err)
			log.Printf("Response was: %+v", response)
			failures = append(failures, err)
		} else {
			log.Printf("Response %d from signer %q passes verification", i, response.SignerID)
		}
	}
	if failed {
		failure := "Errors found during monitoring:"
		for i, fail := range failures {
			failure += fmt.Sprintf("\n%d. %s", i+1, fail.Error())
		}
		return fmt.Errorf(failure)
	}
	log.Println("All signature responses passed, monitoring OK")
	return
}

func makeAuthHeader(req *http.Request, user, token string) string {
	auth := hawk.NewRequestAuth(req,
		&hawk.Credentials{
			ID:   user,
			Key:  token,
			Hash: sha256.New},
		0)
	auth.Ext = fmt.Sprintf("%d", time.Now().Nanosecond())
	payloadhash := auth.PayloadHash("application/json")
	payloadhash.Write([]byte(""))
	auth.SetHash(payloadhash)
	return auth.RequestHeader()
}

// send a message to a predefined sns topic
func sendSoftNotification(id string, format string, a ...interface{}) error {
	if ts, ok := softNotifCache[id]; ok {
		// don't send dup notifications for 24 hours
		if ts.Add(24 * time.Hour).After(time.Now()) {
			log.Printf("silencing soft notification ID %s", id)
			return nil
		}
	}
	if os.Getenv("LAMBDA_TASK_ROOT") == "" || os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS") == "" {
		// We're not running in lambda or the conf isnt ready so don't try to publish to SQS
		log.Printf("soft notification ID %s: %s", id, fmt.Sprintf(format, a...))
		return nil
	}

	svc := sns.New(session.New())
	params := &sns.PublishInput{
		Message:  aws.String(fmt.Sprintf(format, a...)),
		TopicArn: aws.String(os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS")),
	}
	_, err := svc.Publish(params)
	if err != nil {
		return err
	}
	log.Printf("Soft notification send to %q with body: %s", os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS"), *params.Message)
	// add the notification to the cache
	softNotifCache[id] = time.Now()
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"errors"
	"fmt"
	"math/big"
)

// generate creates a random sequence of valid characters of a certain length
func generate(length int, valid charset) (string, error) {
	pass := ""
	errs := []error{}
	for i := 0; i < length; i++ {
		// Attempt to get a random element
		char, err := randElem(valid)
		if err != nil {
			// Record error
			errs = append(errs, err)
			// Tolerate up to 50% error rate, and at least 5
			if len(errs) > max(length/2, 5) {
				return "", tooMany(errs)
			}
		} else {
			// If there is no error, add the rune to the return string
			pass += fmt.Sprintf("%c", char)
		}
	}
	return pass, nil
}

// randElem gets a random rune from a charset
func randElem(set charset) (rune, error) {
	// Create a list to choose a random index from
	list := getList(set)

	if len(list) == 0 {
		return ' ', errors.New("randElem: Character set empty. Check character set generation.")
	}

	// Set the maximum index to choose - casting to big int for crypto/rand
	max := big.NewInt(int64(len(list)))

	// Generate a random index
	// crypto/rand.Int generates a number in [0,max) - See godoc for crypto/rand for more info
	i, err := rand.Int(rand.Reader, max)

	// Error reading from os random source
	if err != nil {
		return ' ', err
	}

	// No error, return rune at random index - casting back to int
	return list[int(i.Int64())], nil
}

// max gets the maximum value of a list of values
func max(is ...int) int {
	m := 0
	for _, i := range is {
		if i > m {
			m = i
		}
	}
	return m
}

// tooMany returns an amalgamated error message listing all of the
// errors that caused the calling function to give up
func tooMany(errs []error) error {
	message := "Too many errors: \n"
	for i, e := range errs {
		message += fmt.Sprintln("\t", i, "-", e)
	}
	return errors.New(message)
}
<commit_msg>Removes punctuation from end of error message<commit_after>package main

import (
	"crypto/rand"
	"errors"
	"fmt"
	"math/big"
)

// generate creates a random sequence of valid characters of a certain length
func generate(length int, valid charset) (string, error) {
	pass := ""
	errs := []error{}
	for i := 0; i < length; i++ {
		// Attempt to get a random element
		char, err := randElem(valid)
		if err != nil {
			// Record error
			errs = append(errs, err)
			// Tolerate up to 50% error rate, and at least 5
			if len(errs) > max(length/2, 5) {
				return "", tooMany(errs)
			}
		} else {
			// If there is no error, add the rune to the return string
			pass += fmt.Sprintf("%c", char)
		}
	}
	return pass, nil
}

// randElem gets a random rune from a charset
func randElem(set charset) (rune, error) {
	// Create a list to choose a random index from
	list := getList(set)

	if len(list) == 0 {
		return ' ', errors.New("randElem: Character set empty. Check character set generation")
	}

	// Set the maximum index to choose - casting to big int for crypto/rand
	max := big.NewInt(int64(len(list)))

	// Generate a random index
	// crypto/rand.Int generates a number in [0,max) - See godoc for crypto/rand for more info
	i, err := rand.Int(rand.Reader, max)

	// Error reading from os random source
	if err != nil {
		return ' ', err
	}

	// No error, return rune at random index - casting back to int
	return list[int(i.Int64())], nil
}

// max gets the maximum value of a list of values
func max(is ...int) int {
	m := 0
	for _, i := range is {
		if i > m {
			m = i
		}
	}
	return m
}

// tooMany returns an amalgamated error message listing all of the
// errors that caused the calling function to give up
func tooMany(errs []error) error {
	message := "Too many errors: \n"
	for i, e := range errs {
		message += fmt.Sprintln("\t", i, "-", e)
	}
	return errors.New(message)
}
<|endoftext|>


<commit_before>package gofakeit

import (
	"strings"
)

// Generate fake information from given string. String should contain {category.subcategory}
//
// Ex: {name.first} - random firstname
//
// Ex: {name.first}###{name.last}@{name.last}.{internet.domain_suffix} - billy834smith@smith.com
//
// Ex: ### - 481 - random numbers
//
// Ex: ??? - fda - random letters
//
// For a complete list possible categories use the Catagories() function.
func Generate(dataVal string) string {
	// Identify items between brackets: {name.first}
	for strings.Count(dataVal, "{") > 0 && strings.Count(dataVal, "}") > 0 {
		catValue := ""
		startIndex := strings.Index(dataVal, "{")
		endIndex := strings.Index(dataVal, "}")
		replace := dataVal[(startIndex + 1):endIndex]
		categories := strings.Split(replace, ".")

		if len(categories) >= 2 && dataCheck([]string{categories[0], categories[1]}) {
			catValue = getRandValue([]string{categories[0], categories[1]})
		}

		dataVal = strings.Replace(dataVal, "{"+replace+"}", catValue, 1)
	}

	// Replace # with numbers
	dataVal = replaceWithNumbers(dataVal)

	// Replace ? with letters
	dataVal = replaceWithLetters(dataVal)

	return dataVal
}
<commit_msg>doc: correct documentation for Generate()<commit_after>package gofakeit

import (
	"strings"
)

// Generate fake information from given string. String should contain {category.subcategory}
//
// Ex: {person.first} - random firstname
//
// Ex: {person.first}###{person.last}@{person.last}.{internet.domain_suffix} - billy834smith@smith.com
//
// Ex: ### - 481 - random numbers
//
// Ex: ??? - fda - random letters
//
// For a complete list possible categories use the Catagories() function.
func Generate(dataVal string) string {
	// Identify items between brackets: {person.first}
	for strings.Count(dataVal, "{") > 0 && strings.Count(dataVal, "}") > 0 {
		catValue := ""
		startIndex := strings.Index(dataVal, "{")
		endIndex := strings.Index(dataVal, "}")
		replace := dataVal[(startIndex + 1):endIndex]
		categories := strings.Split(replace, ".")

		if len(categories) >= 2 && dataCheck([]string{categories[0], categories[1]}) {
			catValue = getRandValue([]string{categories[0], categories[1]})
		}

		dataVal = strings.Replace(dataVal, "{"+replace+"}", catValue, 1)
	}

	// Replace # with numbers
	dataVal = replaceWithNumbers(dataVal)

	// Replace ? with letters
	dataVal = replaceWithLetters(dataVal)

	return dataVal
}
<|endoftext|>


<commit_before>/*
 Copyright 2016 Padduck, LLC

 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 server

import (
	"encoding/json"
	"io"
	"io/ioutil"
	gohttp "net/http"
	"os"
	"path/filepath"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/itsjamie/gin-cors"
	"github.com/pkg/errors"
	ppErrors "github.com/pufferpanel/pufferd/errors"
	"github.com/pufferpanel/apufferi/http"
	"github.com/pufferpanel/pufferd/httphandlers"
	"github.com/pufferpanel/pufferd/logging"
	"github.com/pufferpanel/pufferd/programs"
	"github.com/pufferpanel/pufferd/utils"
	"strconv"
	"strings"
)

var wsupgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *gohttp.Request) bool {
		return true
	},
}

func RegisterRoutes(e *gin.Engine) {
	l := e.Group("/server")
	{
		l.Handle("CONNECT", "/:id/console", func(c *gin.Context) {
			c.Header("Access-Control-Allow-Origin", "*")
			c.Header("Access-Control-Allow-Credentials", "false")
		})
		l.PUT("/:id", httphandlers.OAuth2Handler("server.create", false), CreateServer)
		l.DELETE("/:id", httphandlers.OAuth2Handler("server.delete", true), DeleteServer)
		l.POST("/:id", httphandlers.OAuth2Handler("server.edit", true), EditServer)

		l.GET("/:id/start", httphandlers.OAuth2Handler("server.start", true), StartServer)
		l.GET("/:id/stop", httphandlers.OAuth2Handler("server.stop", true), StopServer)
		l.GET("/:id/kill", httphandlers.OAuth2Handler("server.stop", true), KillServer)

		l.POST("/:id/start", httphandlers.OAuth2Handler("server.start", true), StartServer)
		l.POST("/:id/stop", httphandlers.OAuth2Handler("server.stop", true), StopServer)
		l.POST("/:id/kill", httphandlers.OAuth2Handler("server.stop", true), KillServer)

		l.POST("/:id/install", httphandlers.OAuth2Handler("server.install", true), InstallServer)

		l.GET("/:id/file/*filename", httphandlers.OAuth2Handler("server.file.get", true), GetFile)
		l.PUT("/:id/file/*filename", httphandlers.OAuth2Handler("server.file.put", true), PutFile)
		//l.DELETE("/:id/file/*filename", httphandlers.OAuth2Handler("server.file.delete", true), DeleteFile)
		l.DELETE("/:id/file/*filename", httphandlers.OAuth2Handler("server.file.put", true), DeleteFile)

		l.POST("/:id/console", httphandlers.OAuth2Handler("server.console.send", true), PostConsole)
		l.GET("/:id/console", httphandlers.OAuth2Handler("server.console", true), cors.Middleware(cors.Config{
			Origins:     "*",
			Credentials: true,
		}), GetConsole)
		l.GET("/:id/logs", httphandlers.OAuth2Handler("server.log", true), GetLogs)

		l.GET("/:id/stats", httphandlers.OAuth2Handler("server.stats", true), GetStats)

		l.POST("/:id/reload", httphandlers.OAuth2Handler("server.reload", true), ReloadServer)
	}
	e.GET("/network", httphandlers.OAuth2Handler("server.network", false), NetworkServer)
}

func StartServer(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	server.Start()
	http.Respond(c).Send()
}

func StopServer(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	wait := c.Param("wait")
	if wait == "" || (wait != "true" && wait != "false") {
		wait = "true"
	}

	err := server.Stop()
	if err != nil {
		errorConnection(c, err)
		return
	}

	if wait == "true" {
		err = server.GetEnvironment().WaitForMainProcess()
		if err != nil {
			errorConnection(c, err)
			return
		}
	}
	http.Respond(c).Send()
}

func KillServer(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	err := server.Kill()
	if err != nil {
		errorConnection(c, err)
		return
	}

	http.Respond(c).Send()
}

func CreateServer(c *gin.Context) {
	serverId := c.Param("id")
	prg, _ := programs.Get(serverId)

	if prg != nil {
		http.Respond(c).Status(409).Message("server already exists").Send()
		return
	}

	data := make(map[string]interface{}, 0)
	err := json.NewDecoder(c.Request.Body).Decode(&data)

	if err != nil {
		logging.Error("Error decoding JSON body", err)
		http.Respond(c).Status(400).Message("error parsing json").Data(err).Code(http.MALFORMEDJSON).Send()
		return
	}

	serverType := data["type"].(string)

	if !programs.Create(serverId, serverType, data) {
		errorConnection(c, nil)
	} else {
		http.Respond(c).Send()
	}
}

func DeleteServer(c *gin.Context) {
	item, _ := c.Get("server")
	prg := item.(programs.Program)
	err := programs.Delete(prg.Id())
	if err != nil {
		http.Respond(c).Status(500).Data(err).Message("error deleting server").Send()
	} else {
		http.Respond(c).Send()
	}
}

func InstallServer(c *gin.Context) {
	item, _ := c.Get("server")
	prg := item.(programs.Program)

	http.Respond(c).Send()
	go func() {
		prg.Install()
	}()
}

func EditServer(c *gin.Context) {
	item, _ := c.Get("server")
	prg := item.(programs.Program)

	data := make(map[string]interface{}, 0)
	json.NewDecoder(c.Request.Body).Decode(&data)

	prg.Edit(data)
	http.Respond(c).Send()
}

func GetFile(c *gin.Context) {

	item, _ := c.Get("server")
	server := item.(programs.Program)

	targetPath := c.Param("filename")

	targetFile := utils.JoinPath(server.GetEnvironment().GetRootDirectory(), targetPath)

	if !utils.EnsureAccess(targetFile, server.GetEnvironment().GetRootDirectory()) {
		http.Respond(c).Status(403).Message("invalid file path").Status(http.NOTAUTHORIZED).Send()
		return
	}

	info, err := os.Stat(targetFile)

	if os.IsNotExist(err) {
		errorConnection(c, err)
		return
	}

	if info.IsDir() {
		type FileDesc struct {
			Name      string `json:"name"`
			Modified  int64  `json:"modifyTime"`
			Size      int64  `json:"size,omitempty"`
			File      bool   `json:"isFile"`
			Extension string `json:"extension,omitempty"`
		}

		files, _ := ioutil.ReadDir(targetFile)
		fileNames := make([]interface{}, 0)
		if targetPath != "" && targetPath != "." && targetPath != "/" {
			newFile := &FileDesc{
				Name: "..",
				File: false,
			}
			fileNames = append(fileNames, newFile)
		}
		for _, file := range files {
			newFile := &FileDesc{
				Name: file.Name(),
				File: !file.IsDir(),
			}

			if newFile.File {
				newFile.Size = file.Size()
				newFile.Modified = file.ModTime().Unix()
				newFile.Extension = filepath.Ext(file.Name())
			}

			fileNames = append(fileNames, newFile)
		}
		http.Respond(c).Data(fileNames).Send()
	} else {
		_, err := os.Open(targetFile)
		if err != nil {
			if err == os.ErrNotExist {
				http.Respond(c).Status(404).Code(http.NOFILE).Send()
			} else {
				errorConnection(c, err)
			}
		}
		c.File(targetFile)
	}
}

func PutFile(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	targetPath := c.Param("filename")

	if targetPath == "" {
		c.Status(404)
		return
	}

	targetFile := utils.JoinPath(server.GetEnvironment().GetRootDirectory(), targetPath)

	if !utils.EnsureAccess(targetFile, server.GetEnvironment().GetRootDirectory()) {
		http.Respond(c).Status(403).Message("invalid file path").Status(http.NOTAUTHORIZED).Send()
		return
	}

	_, mkFolder := c.GetQuery("folder")
	if mkFolder {
		err := os.Mkdir(targetFile, 0644)
		if err != nil {
			errorConnection(c, err)
		} else {
			http.Respond(c).Send()
		}
		return
	}
	file, err := os.Create(targetFile)

	if err != nil {
		errorConnection(c, err)
		logging.Error("Error writing file", err)
		return
	}

	_, noform := c.GetQuery("noform")

	var sourceFile io.ReadCloser

	if noform {
		sourceFile = c.Request.Body
	} else {
		c.Request.ParseMultipartForm(32 << 20)
		sourceFile, _, err = c.Request.FormFile("file")
	}

	_, err = io.Copy(file, sourceFile)

	if err != nil {
		errorConnection(c, err)
		logging.Error("Error writing file", err)
	} else {
		http.Respond(c).Send()
	}
}

func DeleteFile(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	targetPath := c.Param("filename")

	targetFile := utils.JoinPath(server.GetEnvironment().GetRootDirectory(), targetPath)

	if !utils.EnsureAccess(targetFile, server.GetEnvironment().GetRootDirectory()) {
		http.Respond(c).Status(403).Message("invalid file path").Status(http.NOTAUTHORIZED).Send()
		return
	}

	err := os.Remove(targetFile)
	if err != nil {
		errorConnection(c, err)
		logging.Error("Failed to delete file", err)
	} else {
		http.Respond(c).Send()
	}
}

func PostConsole(c *gin.Context) {
	item, _ := c.Get("server")
	prg := item.(programs.Program)

	d, _ := ioutil.ReadAll(c.Request.Body)
	cmd := string(d)
	err := prg.Execute(cmd)
	if err != nil {
		errorConnection(c, err)
	} else {
		http.Respond(c).Send()
	}
}

func GetConsole(c *gin.Context) {
	item, _ := c.Get("server")
	program := item.(programs.Program)

	conn, err := wsupgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logging.Error("Error creating websocket", err)
		errorConnection(c, err)
		return
	}
	console, _ := program.GetEnvironment().GetConsole()
	for _, v := range console {
		conn.WriteMessage(websocket.TextMessage, []byte(v))
	}
	program.GetEnvironment().AddListener(conn)
}

func GetStats(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	results, err := server.GetEnvironment().GetStats()
	if err != nil {
		result := make(map[string]interface{})
		result["error"] = err.Error()
		_, isOffline := err.(ppErrors.ServerOffline)
		if isOffline {
			http.Respond(c).Data(result).Status(200).Send()
		} else {
			http.Respond(c).Data(result).Status(500).Send()
		}
	} else {
		http.Respond(c).Data(results).Send()
	}
}

func ReloadServer(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	err := programs.Reload(server.Id())
	if err != nil {
		errorConnection(c, err)
		return
	}
	http.Respond(c).Send()
}

func NetworkServer(c *gin.Context) {

	servers := c.DefaultQuery("ids", "")
	if servers == "" {
		http.Respond(c).Status(400).Code(http.NOSERVERID).Message("no server ids provided").Send()
		return
	}
	serverIds := strings.Split(servers, ",")
	result := make(map[string]string)
	for _, v := range serverIds {
		program, _ := programs.Get(v)
		if program == nil {
			continue
		}
		result[program.Id()] = program.GetNetwork()
	}
	http.Respond(c).Data(result).Send()
}

func GetLogs(c *gin.Context) {
	item, _ := c.Get("server")
	program := item.(programs.Program)

	time := c.DefaultQuery("time", "0")

	castedTime, ok := strconv.ParseInt(time, 10, 64)

	if ok != nil {
		c.AbortWithError(400, errors.New("Time provided is not a valid UNIX time"))
		return
	}

	console, epoch := program.GetEnvironment().GetConsoleFrom(castedTime)
	msg := ""
	for _, k := range console {
		msg += k
	}
	result := make(map[string]interface{})
	result["epoch"] = epoch
	result["logs"] = msg
	http.Respond(c).Data(result).Send()
}

func errorConnection(c *gin.Context, err error) {
	http.Respond(c).Status(500).Code(http.UNKNOWN).Data(err).Message("error handling request").Send()
}
<commit_msg>Fix for file upload issue (#45)<commit_after>/*
 Copyright 2016 Padduck, LLC

 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 server

import (
	"encoding/json"
	"io"
	"io/ioutil"
	"mime"
	gohttp "net/http"
	"os"
	"path/filepath"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/itsjamie/gin-cors"
	"github.com/pkg/errors"
	"github.com/pufferpanel/apufferi/http"
	ppErrors "github.com/pufferpanel/pufferd/errors"
	"github.com/pufferpanel/pufferd/httphandlers"
	"github.com/pufferpanel/pufferd/logging"
	"github.com/pufferpanel/pufferd/programs"
	"github.com/pufferpanel/pufferd/utils"
	"strconv"
	"strings"
)

var wsupgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *gohttp.Request) bool {
		return true
	},
}

func RegisterRoutes(e *gin.Engine) {
	l := e.Group("/server")
	{
		l.Handle("CONNECT", "/:id/console", func(c *gin.Context) {
			c.Header("Access-Control-Allow-Origin", "*")
			c.Header("Access-Control-Allow-Credentials", "false")
		})
		l.PUT("/:id", httphandlers.OAuth2Handler("server.create", false), CreateServer)
		l.DELETE("/:id", httphandlers.OAuth2Handler("server.delete", true), DeleteServer)
		l.POST("/:id", httphandlers.OAuth2Handler("server.edit", true), EditServer)

		l.GET("/:id/start", httphandlers.OAuth2Handler("server.start", true), StartServer)
		l.GET("/:id/stop", httphandlers.OAuth2Handler("server.stop", true), StopServer)
		l.GET("/:id/kill", httphandlers.OAuth2Handler("server.stop", true), KillServer)

		l.POST("/:id/start", httphandlers.OAuth2Handler("server.start", true), StartServer)
		l.POST("/:id/stop", httphandlers.OAuth2Handler("server.stop", true), StopServer)
		l.POST("/:id/kill", httphandlers.OAuth2Handler("server.stop", true), KillServer)

		l.POST("/:id/install", httphandlers.OAuth2Handler("server.install", true), InstallServer)

		l.GET("/:id/file/*filename", httphandlers.OAuth2Handler("server.file.get", true), GetFile)
		l.PUT("/:id/file/*filename", httphandlers.OAuth2Handler("server.file.put", true), PutFile)
		//l.DELETE("/:id/file/*filename", httphandlers.OAuth2Handler("server.file.delete", true), DeleteFile)
		l.DELETE("/:id/file/*filename", httphandlers.OAuth2Handler("server.file.put", true), DeleteFile)

		l.POST("/:id/console", httphandlers.OAuth2Handler("server.console.send", true), PostConsole)
		l.GET("/:id/console", httphandlers.OAuth2Handler("server.console", true), cors.Middleware(cors.Config{
			Origins:     "*",
			Credentials: true,
		}), GetConsole)
		l.GET("/:id/logs", httphandlers.OAuth2Handler("server.log", true), GetLogs)

		l.GET("/:id/stats", httphandlers.OAuth2Handler("server.stats", true), GetStats)

		l.POST("/:id/reload", httphandlers.OAuth2Handler("server.reload", true), ReloadServer)
	}
	e.GET("/network", httphandlers.OAuth2Handler("server.network", false), NetworkServer)
}

func StartServer(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	server.Start()
	http.Respond(c).Send()
}

func StopServer(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	wait := c.Param("wait")
	if wait == "" || (wait != "true" && wait != "false") {
		wait = "true"
	}

	err := server.Stop()
	if err != nil {
		errorConnection(c, err)
		return
	}

	if wait == "true" {
		err = server.GetEnvironment().WaitForMainProcess()
		if err != nil {
			errorConnection(c, err)
			return
		}
	}
	http.Respond(c).Send()
}

func KillServer(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	err := server.Kill()
	if err != nil {
		errorConnection(c, err)
		return
	}

	http.Respond(c).Send()
}

func CreateServer(c *gin.Context) {
	serverId := c.Param("id")
	prg, _ := programs.Get(serverId)

	if prg != nil {
		http.Respond(c).Status(409).Message("server already exists").Send()
		return
	}

	data := make(map[string]interface{}, 0)
	err := json.NewDecoder(c.Request.Body).Decode(&data)

	if err != nil {
		logging.Error("Error decoding JSON body", err)
		http.Respond(c).Status(400).Message("error parsing json").Data(err).Code(http.MALFORMEDJSON).Send()
		return
	}

	serverType := data["type"].(string)

	if !programs.Create(serverId, serverType, data) {
		errorConnection(c, nil)
	} else {
		http.Respond(c).Send()
	}
}

func DeleteServer(c *gin.Context) {
	item, _ := c.Get("server")
	prg := item.(programs.Program)
	err := programs.Delete(prg.Id())
	if err != nil {
		http.Respond(c).Status(500).Data(err).Message("error deleting server").Send()
	} else {
		http.Respond(c).Send()
	}
}

func InstallServer(c *gin.Context) {
	item, _ := c.Get("server")
	prg := item.(programs.Program)

	http.Respond(c).Send()
	go func() {
		prg.Install()
	}()
}

func EditServer(c *gin.Context) {
	item, _ := c.Get("server")
	prg := item.(programs.Program)

	data := make(map[string]interface{}, 0)
	json.NewDecoder(c.Request.Body).Decode(&data)

	prg.Edit(data)
	http.Respond(c).Send()
}

func GetFile(c *gin.Context) {

	item, _ := c.Get("server")
	server := item.(programs.Program)

	targetPath := c.Param("filename")

	targetFile := utils.JoinPath(server.GetEnvironment().GetRootDirectory(), targetPath)

	if !utils.EnsureAccess(targetFile, server.GetEnvironment().GetRootDirectory()) {
		http.Respond(c).Status(403).Message("invalid file path").Status(http.NOTAUTHORIZED).Send()
		return
	}

	info, err := os.Stat(targetFile)

	if os.IsNotExist(err) {
		errorConnection(c, err)
		return
	}

	if info.IsDir() {
		type FileDesc struct {
			Name      string `json:"name"`
			Modified  int64  `json:"modifyTime"`
			Size      int64  `json:"size,omitempty"`
			File      bool   `json:"isFile"`
			Extension string `json:"extension,omitempty"`
		}

		files, _ := ioutil.ReadDir(targetFile)
		fileNames := make([]interface{}, 0)
		if targetPath != "" && targetPath != "." && targetPath != "/" {
			newFile := &FileDesc{
				Name: "..",
				File: false,
			}
			fileNames = append(fileNames, newFile)
		}
		for _, file := range files {
			newFile := &FileDesc{
				Name: file.Name(),
				File: !file.IsDir(),
			}

			if newFile.File {
				newFile.Size = file.Size()
				newFile.Modified = file.ModTime().Unix()
				newFile.Extension = filepath.Ext(file.Name())
			}

			fileNames = append(fileNames, newFile)
		}
		http.Respond(c).Data(fileNames).Send()
	} else {
		_, err := os.Open(targetFile)
		if err != nil {
			if err == os.ErrNotExist {
				http.Respond(c).Status(404).Code(http.NOFILE).Send()
			} else {
				errorConnection(c, err)
			}
		}
		c.File(targetFile)
	}
}

func PutFile(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	targetPath := c.Param("filename")

	if targetPath == "" {
		c.Status(404)
		return
	}

	targetFile := utils.JoinPath(server.GetEnvironment().GetRootDirectory(), targetPath)

	if !utils.EnsureAccess(targetFile, server.GetEnvironment().GetRootDirectory()) {
		http.Respond(c).Status(403).Message("invalid file path").Status(http.NOTAUTHORIZED).Send()
		return
	}

	_, mkFolder := c.GetQuery("folder")
	if mkFolder {
		err := os.Mkdir(targetFile, 0644)
		if err != nil {
			errorConnection(c, err)
		} else {
			http.Respond(c).Send()
		}
		return
	}
	file, err := os.Create(targetFile)
	defer file.Close()

	if err != nil {
		errorConnection(c, err)
		logging.Error("Error writing file", err)
		return
	}

	var sourceFile io.ReadCloser

	v := c.Request.Header.Get("Content-Type")
	if t, _, _ := mime.ParseMediaType(v); t == "multipart/form-data" {
		sourceFile, _, err = c.Request.FormFile("file")
	} else {
		sourceFile = c.Request.Body
	}

	if err != nil {
		errorConnection(c, err)
		logging.Error("Error writing file", err)
	} else {
		_, err = io.Copy(file, sourceFile)
		http.Respond(c).Send()
	}
}

func DeleteFile(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	targetPath := c.Param("filename")

	targetFile := utils.JoinPath(server.GetEnvironment().GetRootDirectory(), targetPath)

	if !utils.EnsureAccess(targetFile, server.GetEnvironment().GetRootDirectory()) {
		http.Respond(c).Status(403).Message("invalid file path").Status(http.NOTAUTHORIZED).Send()
		return
	}

	err := os.Remove(targetFile)
	if err != nil {
		errorConnection(c, err)
		logging.Error("Failed to delete file", err)
	} else {
		http.Respond(c).Send()
	}
}

func PostConsole(c *gin.Context) {
	item, _ := c.Get("server")
	prg := item.(programs.Program)

	d, _ := ioutil.ReadAll(c.Request.Body)
	cmd := string(d)
	err := prg.Execute(cmd)
	if err != nil {
		errorConnection(c, err)
	} else {
		http.Respond(c).Send()
	}
}

func GetConsole(c *gin.Context) {
	item, _ := c.Get("server")
	program := item.(programs.Program)

	conn, err := wsupgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logging.Error("Error creating websocket", err)
		errorConnection(c, err)
		return
	}
	console, _ := program.GetEnvironment().GetConsole()
	for _, v := range console {
		conn.WriteMessage(websocket.TextMessage, []byte(v))
	}
	program.GetEnvironment().AddListener(conn)
}

func GetStats(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	results, err := server.GetEnvironment().GetStats()
	if err != nil {
		result := make(map[string]interface{})
		result["error"] = err.Error()
		_, isOffline := err.(ppErrors.ServerOffline)
		if isOffline {
			http.Respond(c).Data(result).Status(200).Send()
		} else {
			http.Respond(c).Data(result).Status(500).Send()
		}
	} else {
		http.Respond(c).Data(results).Send()
	}
}

func ReloadServer(c *gin.Context) {
	item, _ := c.Get("server")
	server := item.(programs.Program)

	err := programs.Reload(server.Id())
	if err != nil {
		errorConnection(c, err)
		return
	}
	http.Respond(c).Send()
}

func NetworkServer(c *gin.Context) {

	servers := c.DefaultQuery("ids", "")
	if servers == "" {
		http.Respond(c).Status(400).Code(http.NOSERVERID).Message("no server ids provided").Send()
		return
	}
	serverIds := strings.Split(servers, ",")
	result := make(map[string]string)
	for _, v := range serverIds {
		program, _ := programs.Get(v)
		if program == nil {
			continue
		}
		result[program.Id()] = program.GetNetwork()
	}
	http.Respond(c).Data(result).Send()
}

func GetLogs(c *gin.Context) {
	item, _ := c.Get("server")
	program := item.(programs.Program)

	time := c.DefaultQuery("time", "0")

	castedTime, ok := strconv.ParseInt(time, 10, 64)

	if ok != nil {
		c.AbortWithError(400, errors.New("Time provided is not a valid UNIX time"))
		return
	}

	console, epoch := program.GetEnvironment().GetConsoleFrom(castedTime)
	msg := ""
	for _, k := range console {
		msg += k
	}
	result := make(map[string]interface{})
	result["epoch"] = epoch
	result["logs"] = msg
	http.Respond(c).Data(result).Send()
}

func errorConnection(c *gin.Context, err error) {
	http.Respond(c).Status(500).Code(http.UNKNOWN).Data(err).Message("error handling request").Send()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/hawx/serve"
)

type Any interface{}
type Json map[string]Any

func createResponseBody(r *http.Request) string {
	err := r.ParseForm()
	if err != nil {
		log.Println("form:", err)
	}

	form := Json{}
	for k, v := range r.Form {
		form[k] = v
	}

	headers := Json{}
	for k, v := range r.Header {
		headers[k] = strings.Join(v, ", ")
	}

	body, err := ioutil.ReadAll(r.Body)
	defer r.Body.Close()
	if err != nil {
		log.Println("body:", err)
	}

	j := Json{
		"method":  r.Method,
		"url":     r.Host + r.URL.String(),
		"version": r.Proto,
		"headers": headers,
		"body":    string(body),
		"form":    form,
	}

	b, err := json.MarshalIndent(j, "", "  ")
	if err != nil {
		log.Println("json:", err)
		return ""
	}

	return string(b) + "\n"
}

func main() {
	http.HandleFunc("/delay/", func(w http.ResponseWriter, r *http.Request) {
		ms, err := strconv.ParseInt(r.URL.Path[7:], 10, 64)
		if err != nil {
			log.Println(err)
			return
		}

		time.Sleep(time.Duration(ms) * time.Millisecond)
		fmt.Fprintf(w, createResponseBody(r))
	})

	http.HandleFunc("/code/", func(w http.ResponseWriter, r *http.Request) {
		code, err := strconv.ParseInt(r.URL.Path[6:], 10, 0)
		if err != nil {
			log.Println(err)
			return
		}

		w.WriteHeader(int(code))
		fmt.Fprint(w, createResponseBody(r))
	})

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, createResponseBody(r))
	})

	port := flag.String("port", "8080", "")
	socket := flag.String("socket", "", "")

	flag.Parse()

	serve.Serve(*port, *socket, http.DefaultServeMux)
}
<commit_msg>Use hawx/route to handle parameters<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/hawx/route"
	"github.com/hawx/serve"
)

var (
	port   = flag.String("port", "8080", "")
	socket = flag.String("socket", "", "")
)

type Any interface{}
type Json map[string]Any

func createResponseBody(r *http.Request) string {
	err := r.ParseForm()
	if err != nil {
		log.Println("form:", err)
	}

	form := Json{}
	for k, v := range r.Form {
		form[k] = v
	}

	headers := Json{}
	for k, v := range r.Header {
		headers[k] = strings.Join(v, ", ")
	}

	body, err := ioutil.ReadAll(r.Body)
	defer r.Body.Close()
	if err != nil {
		log.Println("body:", err)
	}

	j := Json{
		"method":  r.Method,
		"url":     r.Host + r.URL.String(),
		"version": r.Proto,
		"headers": headers,
		"body":    string(body),
		"form":    form,
	}

	b, err := json.MarshalIndent(j, "", "  ")
	if err != nil {
		log.Println("json:", err)
		return ""
	}

	return string(b) + "\n"
}

func main() {
	flag.Parse()

	route.HandleFunc("/delay/:ms/*path", func(w http.ResponseWriter, r *http.Request) {
		ms, err := strconv.ParseInt(route.Vars(r)["ms"], 10, 64)
		if err != nil {
			log.Println(err)
			return
		}

		time.Sleep(time.Duration(ms) * time.Millisecond)
		fmt.Fprintf(w, createResponseBody(r))
	})

	route.HandleFunc("/code/:code/*path", func(w http.ResponseWriter, r *http.Request) {
		code, err := strconv.ParseInt(route.Vars(r)["code"], 10, 0)
		if err != nil {
			log.Println(err)
			return
		}

		w.WriteHeader(int(code))
		fmt.Fprint(w, createResponseBody(r))
	})

	route.HandleFunc("/*path", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, createResponseBody(r))
	})

	serve.Serve(*port, *socket, route.Default)
}
<|endoftext|>


<commit_before>package response

import (
	"errors"
	"net/http"
	"socialapi/config"
	"socialapi/workers/helper"

	"github.com/koding/bongo"
)

var ErrContentNotFound = errors.New("content not found")

// NewBadRequest is creating a new http response with predifined
// http response properties
func NewBadRequest(err error) (int, http.Header, interface{}, error) {
	if err == nil {
		err = errors.New("request is not valid")
	}

	// make sure errors are outputted
	helper.MustGetLogger().Error("Bad Request: %s", err)

	// do not expose errors to the client
	if config.MustGet().Environment != "dev" {
		err = genericError
	}

	return http.StatusBadRequest, nil, nil, BadRequest{err}
}

// NewAccessDenied sends access denied response back to client
//
// here not to leak info about the resource
// do send NotFound err
func NewAccessDenied(err error) (int, http.Header, interface{}, error) {
	helper.MustGetLogger().Error("Access Denied Err: %s", err.Error())
	return NewNotFound()
}

// HandleResultAndError wraps the function calls and get its reponse,
// assuming the second parameter as error checks it if it is null or not
// if err nor found, returns OK response
func HandleResultAndError(res interface{}, err error) (int, http.Header, interface{}, error) {
	if err == bongo.RecordNotFound {
		return NewNotFound()
	}

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

	return NewOK(res)
}

// HandleResultAndClientError is same as `HandleResultAndError`, but it
// returns the actual error to client as opposed to generic error.
func HandleResultAndClientError(res interface{}, err error) (int, http.Header, interface{}, error) {
	if err != nil {
		return http.StatusBadRequest, nil, nil, err
	}

	return NewOK(res)
}

// NewOK returns http StatusOK response
func NewOK(res interface{}) (int, http.Header, interface{}, error) {
	return http.StatusOK, nil, res, nil
}

// NewNotFound returns http StatusNotFound response
func NewNotFound() (int, http.Header, interface{}, error) {
	return http.StatusNotFound, nil, nil, NotFoundError{ErrContentNotFound}
}

// NewDeleted returns http StatusAccepted response
func NewDeleted() (int, http.Header, interface{}, error) {
	return http.StatusAccepted, nil, nil, nil
}

// NewDefaultOK returns http StatusOK response with `{status:true}` response
func NewDefaultOK() (int, http.Header, interface{}, error) {
	res := map[string]interface{}{
		"status": true,
	}

	return NewOK(res)
}
<commit_msg>wercker: expose socialapi errors<commit_after>package response

import (
	"errors"
	"net/http"
	"socialapi/config"
	"socialapi/workers/helper"

	"github.com/koding/bongo"
)

var ErrContentNotFound = errors.New("content not found")

// NewBadRequest is creating a new http response with predifined
// http response properties
func NewBadRequest(err error) (int, http.Header, interface{}, error) {
	if err == nil {
		err = errors.New("request is not valid")
	}

	// make sure errors are outputted
	helper.MustGetLogger().Error("Bad Request: %s", err)

	// do not expose errors to the client
	env := config.MustGet().Environment
	// do not expose errors to the client.
	if env != "dev" || env != "test" {
		err = genericError
	}

	return http.StatusBadRequest, nil, nil, BadRequest{err}
}

// NewAccessDenied sends access denied response back to client
//
// here not to leak info about the resource
// do send NotFound err
func NewAccessDenied(err error) (int, http.Header, interface{}, error) {
	helper.MustGetLogger().Error("Access Denied Err: %s", err.Error())
	return NewNotFound()
}

// HandleResultAndError wraps the function calls and get its reponse,
// assuming the second parameter as error checks it if it is null or not
// if err nor found, returns OK response
func HandleResultAndError(res interface{}, err error) (int, http.Header, interface{}, error) {
	if err == bongo.RecordNotFound {
		return NewNotFound()
	}

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

	return NewOK(res)
}

// HandleResultAndClientError is same as `HandleResultAndError`, but it
// returns the actual error to client as opposed to generic error.
func HandleResultAndClientError(res interface{}, err error) (int, http.Header, interface{}, error) {
	if err != nil {
		return http.StatusBadRequest, nil, nil, err
	}

	return NewOK(res)
}

// NewOK returns http StatusOK response
func NewOK(res interface{}) (int, http.Header, interface{}, error) {
	return http.StatusOK, nil, res, nil
}

// NewNotFound returns http StatusNotFound response
func NewNotFound() (int, http.Header, interface{}, error) {
	return http.StatusNotFound, nil, nil, NotFoundError{ErrContentNotFound}
}

// NewDeleted returns http StatusAccepted response
func NewDeleted() (int, http.Header, interface{}, error) {
	return http.StatusAccepted, nil, nil, nil
}

// NewDefaultOK returns http StatusOK response with `{status:true}` response
func NewDefaultOK() (int, http.Header, interface{}, error) {
	res := map[string]interface{}{
		"status": true,
	}

	return NewOK(res)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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 vreplication

import (
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"
	"time"

	"context"

	"github.com/golang/protobuf/proto"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/binlog/binlogplayer"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/sqlparser"

	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
	querypb "vitess.io/vitess/go/vt/proto/query"
)

type vcopier struct {
	vr        *vreplicator
	tablePlan *TablePlan
}

func newVCopier(vr *vreplicator) *vcopier {
	return &vcopier{
		vr: vr,
	}
}

// initTablesForCopy (phase 1) identifies the list of tables to be copied and inserts
// them into copy_state. If there are no tables to copy, it explicitly stops
// the stream. Otherwise, the copy phase (phase 2) may think that all tables are copied.
// This will cause us to go into the replication phase (phase 3) without a starting position.
func (vc *vcopier) initTablesForCopy(ctx context.Context) error {
	defer vc.vr.dbClient.Rollback()

	plan, err := buildReplicatorPlan(vc.vr.source.Filter, vc.vr.pkInfoMap, nil)
	if err != nil {
		return err
	}
	if err := vc.vr.dbClient.Begin(); err != nil {
		return err
	}
	// Insert the table list only if at least one table matches.
	if len(plan.TargetTables) != 0 {
		var buf strings.Builder
		buf.WriteString("insert into _vt.copy_state(vrepl_id, table_name) values ")
		prefix := ""
		for name := range plan.TargetTables {
			fmt.Fprintf(&buf, "%s(%d, %s)", prefix, vc.vr.id, encodeString(name))
			prefix = ", "
		}
		if _, err := vc.vr.dbClient.Execute(buf.String()); err != nil {
			return err
		}
		if err := vc.vr.setState(binlogplayer.VReplicationCopying, ""); err != nil {
			return err
		}
	} else {
		if err := vc.vr.setState(binlogplayer.BlpStopped, "There is nothing to replicate"); err != nil {
			return err
		}
	}
	return vc.vr.dbClient.Commit()
}

// copyNext performs a multi-step process on each iteration.
// Step 1: catchup: During this step, it replicates from the source from the last position.
// This is a partial replication: events are applied only to tables or subsets of tables
// that have already been copied. This goes on until replication catches up.
// Step 2: Start streaming. This returns the initial field info along with the GTID
// as of which the snapshot is being streamed.
// Step 3: fastForward: The target is fast-forwarded to the GTID obtained. This should
// be quick because we were mostly caught up as of step 1. This ensures that the
// snapshot of the rows are consistent with the position where the target stopped.
// Step 4: copy rows: Copy the next set of rows from the stream that was started in Step 2.
// This goes on until all rows are copied, or a timeout. In both cases, copyNext
// returns, and the replicator decides whether to invoke copyNext again, or to
// go to the next phase if all the copying is done.
// Steps 2, 3 and 4 are performed by copyTable.
// copyNext also builds the copyState metadata that contains the tables and their last
// primary key that was copied. A nil Result means that nothing has been copied.
// A table that was fully copied is removed from copyState.
func (vc *vcopier) copyNext(ctx context.Context, settings binlogplayer.VRSettings) error {
	qr, err := vc.vr.dbClient.Execute(fmt.Sprintf("select table_name, lastpk from _vt.copy_state where vrepl_id=%d", vc.vr.id))
	if err != nil {
		return err
	}
	var tableToCopy string
	copyState := make(map[string]*sqltypes.Result)
	for _, row := range qr.Rows {
		tableName := row[0].ToString()
		lastpk := row[1].ToString()
		if tableToCopy == "" {
			tableToCopy = tableName
		}
		copyState[tableName] = nil
		if lastpk != "" {
			var r querypb.QueryResult
			if err := proto.UnmarshalText(lastpk, &r); err != nil {
				return err
			}
			copyState[tableName] = sqltypes.Proto3ToResult(&r)
		}
	}
	if len(copyState) == 0 {
		return fmt.Errorf("unexpected: there are no tables to copy")
	}
	if err := vc.catchup(ctx, copyState); err != nil {
		return err
	}
	return vc.copyTable(ctx, tableToCopy, copyState)
}

// catchup replays events to the subset of the tables that have been copied
// until replication is caught up. In order to stop, the seconds behind master has
// to fall below replicationLagTolerance.
func (vc *vcopier) catchup(ctx context.Context, copyState map[string]*sqltypes.Result) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	defer func() {
		vc.vr.stats.PhaseTimings.Record("catchup", time.Now())
	}()

	settings, err := binlogplayer.ReadVRSettings(vc.vr.dbClient, vc.vr.id)
	if err != nil {
		return err
	}
	// If there's no start position, it means we're copying the
	// first table. So, there's nothing to catch up to.
	if settings.StartPos.IsZero() {
		return nil
	}

	// Start vreplication.
	errch := make(chan error, 1)
	go func() {
		errch <- newVPlayer(vc.vr, settings, copyState, mysql.Position{}, "catchup").play(ctx)
	}()

	// Wait for catchup.
	tkr := time.NewTicker(waitRetryTime)
	defer tkr.Stop()
	seconds := int64(replicaLagTolerance / time.Second)
	for {
		sbm := vc.vr.stats.SecondsBehindMaster.Get()
		if sbm < seconds {
			cancel()
			// Make sure vplayer returns before returning.
			<-errch
			return nil
		}
		select {
		case err := <-errch:
			if err != nil {
				return err
			}
			return io.EOF
		case <-ctx.Done():
			// Make sure vplayer returns before returning.
			<-errch
			return io.EOF
		case <-tkr.C:
		}
	}
}

// copyTable performs the synchronized copy of the next set of rows from
// the current table being copied. Each packet received is transactionally
// committed with the lastpk. This allows for consistent resumability.
func (vc *vcopier) copyTable(ctx context.Context, tableName string, copyState map[string]*sqltypes.Result) error {
	defer vc.vr.dbClient.Rollback()
	defer func() {
		vc.vr.stats.PhaseTimings.Record("copy", time.Now())
		vc.vr.stats.CopyLoopCount.Add(1)
	}()

	log.Infof("Copying table %s, lastpk: %v", tableName, copyState[tableName])

	plan, err := buildReplicatorPlan(vc.vr.source.Filter, vc.vr.pkInfoMap, nil)
	if err != nil {
		return err
	}

	initialPlan, ok := plan.TargetTables[tableName]
	if !ok {
		return fmt.Errorf("plan not found for table: %s, current plans are: %#v", tableName, plan.TargetTables)
	}

	ctx, cancel := context.WithTimeout(ctx, copyTimeout)
	defer cancel()

	var lastpkpb *querypb.QueryResult
	if lastpkqr := copyState[tableName]; lastpkqr != nil {
		lastpkpb = sqltypes.ResultToProto3(lastpkqr)
	}

	var pkfields []*querypb.Field
	var updateCopyState *sqlparser.ParsedQuery
	var bv map[string]*querypb.BindVariable
	err = vc.vr.sourceVStreamer.VStreamRows(ctx, initialPlan.SendRule.Filter, lastpkpb, func(rows *binlogdatapb.VStreamRowsResponse) error {
		for {
			select {
			case <-ctx.Done():
				return io.EOF
			default:
			}
			// verify throttler is happy, otherwise keep looping
			if !vc.vr.vre.throttlerClient.ThrottleCheckOKOrWait(ctx) {
				break
			}
		}

		if vc.tablePlan == nil {
			if len(rows.Fields) == 0 {
				return fmt.Errorf("expecting field event first, got: %v", rows)
			}
			if err := vc.fastForward(ctx, copyState, rows.Gtid); err != nil {
				return err
			}
			fieldEvent := &binlogdatapb.FieldEvent{
				TableName: initialPlan.SendRule.Match,
				Fields:    rows.Fields,
			}
			vc.tablePlan, err = plan.buildExecutionPlan(fieldEvent)
			if err != nil {
				return err
			}
			pkfields = rows.Pkfields
			buf := sqlparser.NewTrackedBuffer(nil)
			buf.Myprintf("update _vt.copy_state set lastpk=%a where vrepl_id=%s and table_name=%s", ":lastpk", strconv.Itoa(int(vc.vr.id)), encodeString(tableName))
			updateCopyState = buf.ParsedQuery()
		}
		if len(rows.Rows) == 0 {
			return nil
		}

		// The number of rows we receive depends on the packet size set
		// for the row streamer. Since the packet size is roughly equivalent
		// to data size, this should map to a uniform amount of pages affected
		// per statement. A packet size of 30K will roughly translate to 8
		// mysql pages of 4K each.
		if err := vc.vr.dbClient.Begin(); err != nil {
			return err
		}
		_, err = vc.tablePlan.applyBulkInsert(rows, func(sql string) (*sqltypes.Result, error) {
			start := time.Now()
			qr, err := vc.vr.dbClient.ExecuteWithRetry(ctx, sql)
			if err != nil {
				return nil, err
			}
			vc.vr.stats.QueryTimings.Record("copy", start)

			vc.vr.stats.CopyRowCount.Add(int64(qr.RowsAffected))
			vc.vr.stats.QueryCount.Add("copy", 1)

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

		var buf bytes.Buffer
		err = proto.CompactText(&buf, &querypb.QueryResult{
			Fields: pkfields,
			Rows:   []*querypb.Row{rows.Lastpk},
		})
		if err != nil {
			return err
		}
		bv = map[string]*querypb.BindVariable{
			"lastpk": {
				Type:  sqltypes.VarBinary,
				Value: buf.Bytes(),
			},
		}
		updateState, err := updateCopyState.GenerateQuery(bv, nil)
		if err != nil {
			return err
		}
		if _, err := vc.vr.dbClient.Execute(updateState); err != nil {
			return err
		}

		if err := vc.vr.dbClient.Commit(); err != nil {
			return err
		}
		return nil
	})
	// If there was a timeout, return without an error.
	select {
	case <-ctx.Done():
		log.Infof("Copy of %v stopped at lastpk: %v", tableName, bv)
		return nil
	default:
	}
	if err != nil {
		return err
	}
	log.Infof("Copy of %v finished at lastpk: %v", tableName, bv)
	buf := sqlparser.NewTrackedBuffer(nil)
	buf.Myprintf("delete from _vt.copy_state where vrepl_id=%s and table_name=%s", strconv.Itoa(int(vc.vr.id)), encodeString(tableName))
	if _, err := vc.vr.dbClient.Execute(buf.String()); err != nil {
		return err
	}
	return nil
}

func (vc *vcopier) fastForward(ctx context.Context, copyState map[string]*sqltypes.Result, gtid string) error {
	defer func() {
		vc.vr.stats.PhaseTimings.Record("fastforward", time.Now())
	}()
	pos, err := mysql.DecodePosition(gtid)
	if err != nil {
		return err
	}
	settings, err := binlogplayer.ReadVRSettings(vc.vr.dbClient, vc.vr.id)
	if err != nil {
		return err
	}
	if settings.StartPos.IsZero() {
		update := binlogplayer.GenerateUpdatePos(vc.vr.id, pos, time.Now().Unix(), 0)
		_, err := vc.vr.dbClient.Execute(update)
		return err
	}
	return newVPlayer(vc.vr, settings, copyState, pos, "fastforward").play(ctx)
}
<commit_msg>fixed throttle logic in vcopier<commit_after>/*
Copyright 2019 The Vitess Authors.

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 vreplication

import (
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"
	"time"

	"context"

	"github.com/golang/protobuf/proto"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/binlog/binlogplayer"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/sqlparser"

	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
	querypb "vitess.io/vitess/go/vt/proto/query"
)

type vcopier struct {
	vr        *vreplicator
	tablePlan *TablePlan
}

func newVCopier(vr *vreplicator) *vcopier {
	return &vcopier{
		vr: vr,
	}
}

// initTablesForCopy (phase 1) identifies the list of tables to be copied and inserts
// them into copy_state. If there are no tables to copy, it explicitly stops
// the stream. Otherwise, the copy phase (phase 2) may think that all tables are copied.
// This will cause us to go into the replication phase (phase 3) without a starting position.
func (vc *vcopier) initTablesForCopy(ctx context.Context) error {
	defer vc.vr.dbClient.Rollback()

	plan, err := buildReplicatorPlan(vc.vr.source.Filter, vc.vr.pkInfoMap, nil)
	if err != nil {
		return err
	}
	if err := vc.vr.dbClient.Begin(); err != nil {
		return err
	}
	// Insert the table list only if at least one table matches.
	if len(plan.TargetTables) != 0 {
		var buf strings.Builder
		buf.WriteString("insert into _vt.copy_state(vrepl_id, table_name) values ")
		prefix := ""
		for name := range plan.TargetTables {
			fmt.Fprintf(&buf, "%s(%d, %s)", prefix, vc.vr.id, encodeString(name))
			prefix = ", "
		}
		if _, err := vc.vr.dbClient.Execute(buf.String()); err != nil {
			return err
		}
		if err := vc.vr.setState(binlogplayer.VReplicationCopying, ""); err != nil {
			return err
		}
	} else {
		if err := vc.vr.setState(binlogplayer.BlpStopped, "There is nothing to replicate"); err != nil {
			return err
		}
	}
	return vc.vr.dbClient.Commit()
}

// copyNext performs a multi-step process on each iteration.
// Step 1: catchup: During this step, it replicates from the source from the last position.
// This is a partial replication: events are applied only to tables or subsets of tables
// that have already been copied. This goes on until replication catches up.
// Step 2: Start streaming. This returns the initial field info along with the GTID
// as of which the snapshot is being streamed.
// Step 3: fastForward: The target is fast-forwarded to the GTID obtained. This should
// be quick because we were mostly caught up as of step 1. This ensures that the
// snapshot of the rows are consistent with the position where the target stopped.
// Step 4: copy rows: Copy the next set of rows from the stream that was started in Step 2.
// This goes on until all rows are copied, or a timeout. In both cases, copyNext
// returns, and the replicator decides whether to invoke copyNext again, or to
// go to the next phase if all the copying is done.
// Steps 2, 3 and 4 are performed by copyTable.
// copyNext also builds the copyState metadata that contains the tables and their last
// primary key that was copied. A nil Result means that nothing has been copied.
// A table that was fully copied is removed from copyState.
func (vc *vcopier) copyNext(ctx context.Context, settings binlogplayer.VRSettings) error {
	qr, err := vc.vr.dbClient.Execute(fmt.Sprintf("select table_name, lastpk from _vt.copy_state where vrepl_id=%d", vc.vr.id))
	if err != nil {
		return err
	}
	var tableToCopy string
	copyState := make(map[string]*sqltypes.Result)
	for _, row := range qr.Rows {
		tableName := row[0].ToString()
		lastpk := row[1].ToString()
		if tableToCopy == "" {
			tableToCopy = tableName
		}
		copyState[tableName] = nil
		if lastpk != "" {
			var r querypb.QueryResult
			if err := proto.UnmarshalText(lastpk, &r); err != nil {
				return err
			}
			copyState[tableName] = sqltypes.Proto3ToResult(&r)
		}
	}
	if len(copyState) == 0 {
		return fmt.Errorf("unexpected: there are no tables to copy")
	}
	if err := vc.catchup(ctx, copyState); err != nil {
		return err
	}
	return vc.copyTable(ctx, tableToCopy, copyState)
}

// catchup replays events to the subset of the tables that have been copied
// until replication is caught up. In order to stop, the seconds behind master has
// to fall below replicationLagTolerance.
func (vc *vcopier) catchup(ctx context.Context, copyState map[string]*sqltypes.Result) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	defer func() {
		vc.vr.stats.PhaseTimings.Record("catchup", time.Now())
	}()

	settings, err := binlogplayer.ReadVRSettings(vc.vr.dbClient, vc.vr.id)
	if err != nil {
		return err
	}
	// If there's no start position, it means we're copying the
	// first table. So, there's nothing to catch up to.
	if settings.StartPos.IsZero() {
		return nil
	}

	// Start vreplication.
	errch := make(chan error, 1)
	go func() {
		errch <- newVPlayer(vc.vr, settings, copyState, mysql.Position{}, "catchup").play(ctx)
	}()

	// Wait for catchup.
	tkr := time.NewTicker(waitRetryTime)
	defer tkr.Stop()
	seconds := int64(replicaLagTolerance / time.Second)
	for {
		sbm := vc.vr.stats.SecondsBehindMaster.Get()
		if sbm < seconds {
			cancel()
			// Make sure vplayer returns before returning.
			<-errch
			return nil
		}
		select {
		case err := <-errch:
			if err != nil {
				return err
			}
			return io.EOF
		case <-ctx.Done():
			// Make sure vplayer returns before returning.
			<-errch
			return io.EOF
		case <-tkr.C:
		}
	}
}

// copyTable performs the synchronized copy of the next set of rows from
// the current table being copied. Each packet received is transactionally
// committed with the lastpk. This allows for consistent resumability.
func (vc *vcopier) copyTable(ctx context.Context, tableName string, copyState map[string]*sqltypes.Result) error {
	defer vc.vr.dbClient.Rollback()
	defer func() {
		vc.vr.stats.PhaseTimings.Record("copy", time.Now())
		vc.vr.stats.CopyLoopCount.Add(1)
	}()

	log.Infof("Copying table %s, lastpk: %v", tableName, copyState[tableName])

	plan, err := buildReplicatorPlan(vc.vr.source.Filter, vc.vr.pkInfoMap, nil)
	if err != nil {
		return err
	}

	initialPlan, ok := plan.TargetTables[tableName]
	if !ok {
		return fmt.Errorf("plan not found for table: %s, current plans are: %#v", tableName, plan.TargetTables)
	}

	ctx, cancel := context.WithTimeout(ctx, copyTimeout)
	defer cancel()

	var lastpkpb *querypb.QueryResult
	if lastpkqr := copyState[tableName]; lastpkqr != nil {
		lastpkpb = sqltypes.ResultToProto3(lastpkqr)
	}

	var pkfields []*querypb.Field
	var updateCopyState *sqlparser.ParsedQuery
	var bv map[string]*querypb.BindVariable
	err = vc.vr.sourceVStreamer.VStreamRows(ctx, initialPlan.SendRule.Filter, lastpkpb, func(rows *binlogdatapb.VStreamRowsResponse) error {
		for {
			select {
			case <-ctx.Done():
				return io.EOF
			default:
			}
			// verify throttler is happy, otherwise keep looping
			if vc.vr.vre.throttlerClient.ThrottleCheckOKOrWait(ctx) {
				break
			}
		}

		if vc.tablePlan == nil {
			if len(rows.Fields) == 0 {
				return fmt.Errorf("expecting field event first, got: %v", rows)
			}
			if err := vc.fastForward(ctx, copyState, rows.Gtid); err != nil {
				return err
			}
			fieldEvent := &binlogdatapb.FieldEvent{
				TableName: initialPlan.SendRule.Match,
				Fields:    rows.Fields,
			}
			vc.tablePlan, err = plan.buildExecutionPlan(fieldEvent)
			if err != nil {
				return err
			}
			pkfields = rows.Pkfields
			buf := sqlparser.NewTrackedBuffer(nil)
			buf.Myprintf("update _vt.copy_state set lastpk=%a where vrepl_id=%s and table_name=%s", ":lastpk", strconv.Itoa(int(vc.vr.id)), encodeString(tableName))
			updateCopyState = buf.ParsedQuery()
		}
		if len(rows.Rows) == 0 {
			return nil
		}

		// The number of rows we receive depends on the packet size set
		// for the row streamer. Since the packet size is roughly equivalent
		// to data size, this should map to a uniform amount of pages affected
		// per statement. A packet size of 30K will roughly translate to 8
		// mysql pages of 4K each.
		if err := vc.vr.dbClient.Begin(); err != nil {
			return err
		}
		_, err = vc.tablePlan.applyBulkInsert(rows, func(sql string) (*sqltypes.Result, error) {
			start := time.Now()
			qr, err := vc.vr.dbClient.ExecuteWithRetry(ctx, sql)
			if err != nil {
				return nil, err
			}
			vc.vr.stats.QueryTimings.Record("copy", start)

			vc.vr.stats.CopyRowCount.Add(int64(qr.RowsAffected))
			vc.vr.stats.QueryCount.Add("copy", 1)

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

		var buf bytes.Buffer
		err = proto.CompactText(&buf, &querypb.QueryResult{
			Fields: pkfields,
			Rows:   []*querypb.Row{rows.Lastpk},
		})
		if err != nil {
			return err
		}
		bv = map[string]*querypb.BindVariable{
			"lastpk": {
				Type:  sqltypes.VarBinary,
				Value: buf.Bytes(),
			},
		}
		updateState, err := updateCopyState.GenerateQuery(bv, nil)
		if err != nil {
			return err
		}
		if _, err := vc.vr.dbClient.Execute(updateState); err != nil {
			return err
		}

		if err := vc.vr.dbClient.Commit(); err != nil {
			return err
		}
		return nil
	})
	// If there was a timeout, return without an error.
	select {
	case <-ctx.Done():
		log.Infof("Copy of %v stopped at lastpk: %v", tableName, bv)
		return nil
	default:
	}
	if err != nil {
		return err
	}
	log.Infof("Copy of %v finished at lastpk: %v", tableName, bv)
	buf := sqlparser.NewTrackedBuffer(nil)
	buf.Myprintf("delete from _vt.copy_state where vrepl_id=%s and table_name=%s", strconv.Itoa(int(vc.vr.id)), encodeString(tableName))
	if _, err := vc.vr.dbClient.Execute(buf.String()); err != nil {
		return err
	}
	return nil
}

func (vc *vcopier) fastForward(ctx context.Context, copyState map[string]*sqltypes.Result, gtid string) error {
	defer func() {
		vc.vr.stats.PhaseTimings.Record("fastforward", time.Now())
	}()
	pos, err := mysql.DecodePosition(gtid)
	if err != nil {
		return err
	}
	settings, err := binlogplayer.ReadVRSettings(vc.vr.dbClient, vc.vr.id)
	if err != nil {
		return err
	}
	if settings.StartPos.IsZero() {
		update := binlogplayer.GenerateUpdatePos(vc.vr.id, pos, time.Now().Unix(), 0)
		_, err := vc.vr.dbClient.Execute(update)
		return err
	}
	return newVPlayer(vc.vr, settings, copyState, pos, "fastforward").play(ctx)
}
<|endoftext|>


<commit_before>package ole

import (
	"errors"
	"fmt"
	"time"

	"github.com/go-ole/go-ole"
	"github.com/go-ole/go-ole/oleutil"
	"github.com/yuin/gopher-lua"
)

type Lua = *lua.LState

var initializedRequired = true

type capsuleT struct {
	Data *ole.IDispatch
}

type methodT struct {
	Name string
	Data *ole.IDispatch
}

func (c capsuleT) ToLValue(L Lua) lua.LValue {
	ud := L.NewUserData()
	ud.Value = &c
	meta := L.NewTable()
	L.SetField(meta, "__gc", L.NewFunction(gc))
	L.SetField(meta, "__index", L.NewFunction(index))
	L.SetField(meta, "__newindex", L.NewFunction(set))
	L.SetMetatable(ud, meta)
	return ud
}

func gc(L Lua) int {
	ud := L.ToUserData(1)
	p, ok := ud.Value.(*capsuleT)
	if !ok {
		return lerror(L, "gc: not capsult_t instance")
	}
	if p.Data != nil {
		p.Data.Release()
		p.Data = nil
	}
	return 0
}

func lua2interface(L Lua, index int) (interface{}, error) {
	valueTmp := L.Get(index)
	if valueTmp == lua.LNil {
		return nil, nil
	} else if valueTmp == lua.LTrue {
		return true, nil
	} else if valueTmp == lua.LFalse {
		return false, nil
	}
	switch value := valueTmp.(type) {
	default:
		return nil, errors.New("lua2interface: not support type")
	case lua.LString:
		return string(value), nil
	case lua.LNumber:
		return int(value), nil
	case *lua.LUserData:
		c, ok := value.Value.(*capsuleT)
		if !ok {
			return nil, errors.New("lua2interface: not a OBJECT")
		}
		return c.Data, nil
	}
}

func lua2interfaceS(L Lua, start, end int) ([]interface{}, error) {
	result := make([]interface{}, end-start+1)
	for i := start; i <= end; i++ {
		val, err := lua2interface(L, i)
		if err != nil {
			return nil, err
		}
		result[i-start] = val
	}
	return result, nil
}

// this:_call("METHODNAME",params...)
func call1(L Lua) int {
	ud, ok := L.Get(1).(*lua.LUserData)
	if !ok { // OBJECT_T
		return lerror(L, "call1: not found object")
	}
	p, ok := ud.Value.(*capsuleT)
	if !ok {
		return lerror(L, "call1: not found capsuleT")
	}
	name, ok := L.Get(2).(lua.LString)
	if !ok {
		return lerror(L, "call1: not found methodname")
	}
	return callCommon(L, p.Data, string(name))
}

// this:METHODNAME(params...)
func call2(L Lua) int {
	ud, ok := L.Get(1).(*lua.LUserData)
	if !ok {
		return lerror(L, "call2: not found userdata for methodT")
	}
	method, ok := ud.Value.(*methodT)
	if !ok || method.Name == "" {
		return lerror(L, "call2: not found methodT")
	}
	ud, ok = L.Get(2).(*lua.LUserData)
	if !ok {
		return lerror(L, "call2: not found userdata for object_t")
	}
	obj, ok := ud.Value.(*capsuleT)
	if !ok {
		if method.Data == nil {
			return lerror(L, "call2: receiver is not found")
		}
		return callCommon(L, method.Data, method.Name)
		// this code enables `OLEOBJ.PROPERTY.PROPERTY:METHOD()`
	}
	if obj.Data == nil {
		return lerror(L, "call2: OLEOBJECT(): the receiver is null")
	}
	return callCommon(L, obj.Data, method.Name)
}

func callCommon(L Lua, com1 *ole.IDispatch, name string) int {
	count := L.GetTop()
	params, err := lua2interfaceS(L, 3, count)
	if err != nil {
		return lerror(L, fmt.Sprintf("callCommon: %s", err.Error()))
	}
	result, err := com1.CallMethod(name, params...)
	if err != nil {
		return lerror(L, fmt.Sprintf("oleutil.CallMethod(%s): %s", name, err.Error()))
	}
	L.Push(variantToLValue(L, result))
	return 1
}

func set(L Lua) int {
	ud, ok := L.Get(1).(*lua.LUserData)
	if !ok {
		return lerror(L, "set: the 1st argument is not usedata")
	}
	p, ok := ud.Value.(*capsuleT)
	if !ok {
		return lerror(L, "set: the 1st argument is not *capsuleT")
	}
	name, ok := L.Get(2).(lua.LString)
	if !ok {
		return lerror(L, "set: the 2nd argument is not string")
	}
	key, err := lua2interfaceS(L, 3, L.GetTop())
	if err != nil {
		return lerror(L, fmt.Sprintf("set: %s", err.Error()))
	}
	p.Data.PutProperty(string(name), key...)
	L.Push(lua.LTrue)
	L.Push(lua.LNil)
	return 2
}

func get(L Lua) int {
	ud, ok := L.Get(1).(*lua.LUserData)
	if !ok {
		return lerror(L, "get: 1st argument is not a userdata.")
	}
	p, ok := ud.Value.(*capsuleT)
	if !ok {
		return lerror(L, "get: 1st argument is not *capsuleT")
	}

	name, ok := L.Get(2).(lua.LString)
	if !ok {
		return lerror(L, "get: 2nd argument is not string")
	}

	key, err := lua2interfaceS(L, 3, L.GetTop())
	if err != nil {
		return lerror(L, fmt.Sprintf("get: %s", err.Error()))
	}
	result, err := p.Data.GetProperty(string(name), key...)
	if err != nil {
		return lerror(L, fmt.Sprintf("oleutil.GetProperty: %s", err.Error()))
	}
	L.Push(variantToLValue(L, result))
	return 1
}

func indexSub(L Lua, thisIndex int, nameIndex int) int {
	name, ok := L.Get(nameIndex).(lua.LString)
	if !ok {
		return lerror(L, "indexSub: not a string")
	}
	switch string(name) {
	case "_call":
		L.Push(L.NewFunction(call1))
		L.Push(lua.LNil)
		return 2
	case "_set":
		L.Push(L.NewFunction(set))
		L.Push(lua.LNil)
		return 2
	case "_get":
		L.Push(L.NewFunction(get))
		L.Push(lua.LNil)
		return 2
	default:
		m := &methodT{Name: string(name)}
		if ud, ok := L.Get(thisIndex).(*lua.LUserData); ok {
			if p, ok := ud.Value.(*capsuleT); ok {
				m.Data = p.Data
			}
		}
		ud := L.NewUserData()
		ud.Value = m

		meta := L.NewTable()
		L.SetField(meta, "__newindex", L.NewFunction(set))
		L.SetField(meta, "__call", L.NewFunction(call2))
		L.SetField(meta, "__index", L.NewFunction(get2))
		L.SetMetatable(ud, meta)
		L.Push(ud)

		return 1
	}
}

func index(L Lua) int {
	return indexSub(L, 1, 2)
}

// THIS.member.member
func get2(L Lua) int {
	ud, ok := L.Get(1).(*lua.LUserData)
	if !ok {
		return lerror(L, "get2: not a userdata")
	}
	m, ok := ud.Value.(*methodT)
	if !ok {
		return lerror(L, "get: not a methodT")
	}
	result, err := m.Data.GetProperty(m.Name)
	if err != nil {
		return lerror(L, fmt.Sprintf("oleutil.GetProperty: %s", err.Error()))
	}
	L.Push(variantToLValue(L, result))
	return indexSub(L, 3, 2)
}

// CreateObject creates Lua-Object to access COM
func CreateObject(L Lua) int {
	if initializedRequired {
		ole.CoInitialize(0)
		initializedRequired = false
	}
	name, ok := L.Get(1).(lua.LString)
	if !ok {
		return lerror(L, "CreateObject: parameter not a string")
	}
	unknown, err := oleutil.CreateObject(string(name))
	if err != nil {
		return lerror(L, fmt.Sprintf("oleutil.CreateObject: %s", err.Error()))
	}
	obj, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return lerror(L, fmt.Sprintf("unknown.QueryInterfce: %s", err.Error()))
	}
	L.Push(capsuleT{obj}.ToLValue(L))
	return 1
}

func lerror(L Lua, s string) int {
	L.Push(lua.LNil)
	L.Push(lua.LString(s))
	return 2
}

func variantToLValue(L Lua, v *ole.VARIANT) lua.LValue {
	switch v.VT {
	case ole.VT_I1:
		return lua.LNumber(v.Value().(int))
	case ole.VT_UI1:
		return lua.LNumber(v.Value().(uint8))
	case ole.VT_I2:
		return lua.LNumber(v.Value().(int16))
	case ole.VT_UI2:
		return lua.LNumber(v.Value().(uint16))
	case ole.VT_I4:
		return lua.LNumber(v.Value().(int32))
	case ole.VT_UI4:
		return lua.LNumber(v.Value().(uint32))
	case ole.VT_I8:
		return lua.LNumber(v.Value().(int64))
	case ole.VT_UI8:
		return lua.LNumber(v.Value().(uint64))
	case ole.VT_INT:
		return lua.LNumber(v.Value().(int))
	case ole.VT_UINT:
		return lua.LNumber(v.Value().(uint))
	case ole.VT_INT_PTR:
		return lua.LNumber(v.Value().(uintptr))
	case ole.VT_UINT_PTR:
		return lua.LNumber(v.Value().(uintptr))
	case ole.VT_R4:
		return lua.LNumber(v.Value().(float32))
	case ole.VT_R8:
		return lua.LNumber(v.Value().(float64))
	case ole.VT_BSTR:
		return lua.LString(v.ToString())
	case ole.VT_DATE:
		if date, ok := v.Value().(time.Time); ok {
			t := L.NewTable()
			L.SetField(t, "year", lua.LNumber(date.Year()))
			L.SetField(t, "month", lua.LNumber(int(date.Month())))
			L.SetField(t, "day", lua.LNumber(date.Day()))
			L.SetField(t, "hour", lua.LNumber(date.Hour()))
			L.SetField(t, "min", lua.LNumber(date.Minute()))
			L.SetField(t, "sec", lua.LNumber(date.Second()))
			return t
		} else if floatValue, ok := v.Value().(float64); ok {
			return lua.LNumber(floatValue)
		} else {
			panic("can not convert ole.VT_DATE")
		}
	case ole.VT_UNKNOWN:
		panic("can not convert ole.VT_UNKNOW")
	case ole.VT_DISPATCH:
		return capsuleT{v.ToIDispatch()}.ToLValue(L)
	case ole.VT_BOOL:
		if v.Value().(bool) {
			return lua.LTrue
		} else {
			return lua.LFalse
		}
	default:
		return lua.LNil
	}
}
<commit_msg>Fix: lua.LNumber was treated as integer. It should be as float64<commit_after>package ole

import (
	"errors"
	"fmt"
	"time"

	"github.com/go-ole/go-ole"
	"github.com/go-ole/go-ole/oleutil"
	"github.com/yuin/gopher-lua"
)

type Lua = *lua.LState

var initializedRequired = true

type capsuleT struct {
	Data *ole.IDispatch
}

type methodT struct {
	Name string
	Data *ole.IDispatch
}

func (c capsuleT) ToLValue(L Lua) lua.LValue {
	ud := L.NewUserData()
	ud.Value = &c
	meta := L.NewTable()
	L.SetField(meta, "__gc", L.NewFunction(gc))
	L.SetField(meta, "__index", L.NewFunction(index))
	L.SetField(meta, "__newindex", L.NewFunction(set))
	L.SetMetatable(ud, meta)
	return ud
}

func gc(L Lua) int {
	ud := L.ToUserData(1)
	p, ok := ud.Value.(*capsuleT)
	if !ok {
		return lerror(L, "gc: not capsult_t instance")
	}
	if p.Data != nil {
		p.Data.Release()
		p.Data = nil
	}
	return 0
}

func lua2interface(L Lua, index int) (interface{}, error) {
	valueTmp := L.Get(index)
	if valueTmp == lua.LNil {
		return nil, nil
	} else if valueTmp == lua.LTrue {
		return true, nil
	} else if valueTmp == lua.LFalse {
		return false, nil
	}
	switch value := valueTmp.(type) {
	default:
		return nil, errors.New("lua2interface: not support type")
	case lua.LString:
		return string(value), nil
	case lua.LNumber:
		return float64(value), nil
	case *lua.LUserData:
		c, ok := value.Value.(*capsuleT)
		if !ok {
			return nil, errors.New("lua2interface: not a OBJECT")
		}
		return c.Data, nil
	}
}

func lua2interfaceS(L Lua, start, end int) ([]interface{}, error) {
	result := make([]interface{}, end-start+1)
	for i := start; i <= end; i++ {
		val, err := lua2interface(L, i)
		if err != nil {
			return nil, err
		}
		result[i-start] = val
	}
	return result, nil
}

// this:_call("METHODNAME",params...)
func call1(L Lua) int {
	ud, ok := L.Get(1).(*lua.LUserData)
	if !ok { // OBJECT_T
		return lerror(L, "call1: not found object")
	}
	p, ok := ud.Value.(*capsuleT)
	if !ok {
		return lerror(L, "call1: not found capsuleT")
	}
	name, ok := L.Get(2).(lua.LString)
	if !ok {
		return lerror(L, "call1: not found methodname")
	}
	return callCommon(L, p.Data, string(name))
}

// this:METHODNAME(params...)
func call2(L Lua) int {
	ud, ok := L.Get(1).(*lua.LUserData)
	if !ok {
		return lerror(L, "call2: not found userdata for methodT")
	}
	method, ok := ud.Value.(*methodT)
	if !ok || method.Name == "" {
		return lerror(L, "call2: not found methodT")
	}
	ud, ok = L.Get(2).(*lua.LUserData)
	if !ok {
		return lerror(L, "call2: not found userdata for object_t")
	}
	obj, ok := ud.Value.(*capsuleT)
	if !ok {
		if method.Data == nil {
			return lerror(L, "call2: receiver is not found")
		}
		return callCommon(L, method.Data, method.Name)
		// this code enables `OLEOBJ.PROPERTY.PROPERTY:METHOD()`
	}
	if obj.Data == nil {
		return lerror(L, "call2: OLEOBJECT(): the receiver is null")
	}
	return callCommon(L, obj.Data, method.Name)
}

func callCommon(L Lua, com1 *ole.IDispatch, name string) int {
	count := L.GetTop()
	params, err := lua2interfaceS(L, 3, count)
	if err != nil {
		return lerror(L, fmt.Sprintf("callCommon: %s", err.Error()))
	}
	result, err := com1.CallMethod(name, params...)
	if err != nil {
		return lerror(L, fmt.Sprintf("oleutil.CallMethod(%s): %s", name, err.Error()))
	}
	L.Push(variantToLValue(L, result))
	return 1
}

func set(L Lua) int {
	ud, ok := L.Get(1).(*lua.LUserData)
	if !ok {
		return lerror(L, "set: the 1st argument is not usedata")
	}
	p, ok := ud.Value.(*capsuleT)
	if !ok {
		return lerror(L, "set: the 1st argument is not *capsuleT")
	}
	name, ok := L.Get(2).(lua.LString)
	if !ok {
		return lerror(L, "set: the 2nd argument is not string")
	}
	key, err := lua2interfaceS(L, 3, L.GetTop())
	if err != nil {
		return lerror(L, fmt.Sprintf("set: %s", err.Error()))
	}
	p.Data.PutProperty(string(name), key...)
	L.Push(lua.LTrue)
	L.Push(lua.LNil)
	return 2
}

func get(L Lua) int {
	ud, ok := L.Get(1).(*lua.LUserData)
	if !ok {
		return lerror(L, "get: 1st argument is not a userdata.")
	}
	p, ok := ud.Value.(*capsuleT)
	if !ok {
		return lerror(L, "get: 1st argument is not *capsuleT")
	}

	name, ok := L.Get(2).(lua.LString)
	if !ok {
		return lerror(L, "get: 2nd argument is not string")
	}

	key, err := lua2interfaceS(L, 3, L.GetTop())
	if err != nil {
		return lerror(L, fmt.Sprintf("get: %s", err.Error()))
	}
	result, err := p.Data.GetProperty(string(name), key...)
	if err != nil {
		return lerror(L, fmt.Sprintf("oleutil.GetProperty: %s", err.Error()))
	}
	L.Push(variantToLValue(L, result))
	return 1
}

func indexSub(L Lua, thisIndex int, nameIndex int) int {
	name, ok := L.Get(nameIndex).(lua.LString)
	if !ok {
		return lerror(L, "indexSub: not a string")
	}
	switch string(name) {
	case "_call":
		L.Push(L.NewFunction(call1))
		L.Push(lua.LNil)
		return 2
	case "_set":
		L.Push(L.NewFunction(set))
		L.Push(lua.LNil)
		return 2
	case "_get":
		L.Push(L.NewFunction(get))
		L.Push(lua.LNil)
		return 2
	default:
		m := &methodT{Name: string(name)}
		if ud, ok := L.Get(thisIndex).(*lua.LUserData); ok {
			if p, ok := ud.Value.(*capsuleT); ok {
				m.Data = p.Data
			}
		}
		ud := L.NewUserData()
		ud.Value = m

		meta := L.NewTable()
		L.SetField(meta, "__newindex", L.NewFunction(set))
		L.SetField(meta, "__call", L.NewFunction(call2))
		L.SetField(meta, "__index", L.NewFunction(get2))
		L.SetMetatable(ud, meta)
		L.Push(ud)

		return 1
	}
}

func index(L Lua) int {
	return indexSub(L, 1, 2)
}

// THIS.member.member
func get2(L Lua) int {
	ud, ok := L.Get(1).(*lua.LUserData)
	if !ok {
		return lerror(L, "get2: not a userdata")
	}
	m, ok := ud.Value.(*methodT)
	if !ok {
		return lerror(L, "get: not a methodT")
	}
	result, err := m.Data.GetProperty(m.Name)
	if err != nil {
		return lerror(L, fmt.Sprintf("oleutil.GetProperty: %s", err.Error()))
	}
	L.Push(variantToLValue(L, result))
	return indexSub(L, 3, 2)
}

// CreateObject creates Lua-Object to access COM
func CreateObject(L Lua) int {
	if initializedRequired {
		ole.CoInitialize(0)
		initializedRequired = false
	}
	name, ok := L.Get(1).(lua.LString)
	if !ok {
		return lerror(L, "CreateObject: parameter not a string")
	}
	unknown, err := oleutil.CreateObject(string(name))
	if err != nil {
		return lerror(L, fmt.Sprintf("oleutil.CreateObject: %s", err.Error()))
	}
	obj, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return lerror(L, fmt.Sprintf("unknown.QueryInterfce: %s", err.Error()))
	}
	L.Push(capsuleT{obj}.ToLValue(L))
	return 1
}

func lerror(L Lua, s string) int {
	L.Push(lua.LNil)
	L.Push(lua.LString(s))
	return 2
}

func variantToLValue(L Lua, v *ole.VARIANT) lua.LValue {
	switch v.VT {
	case ole.VT_I1:
		return lua.LNumber(v.Value().(int))
	case ole.VT_UI1:
		return lua.LNumber(v.Value().(uint8))
	case ole.VT_I2:
		return lua.LNumber(v.Value().(int16))
	case ole.VT_UI2:
		return lua.LNumber(v.Value().(uint16))
	case ole.VT_I4:
		return lua.LNumber(v.Value().(int32))
	case ole.VT_UI4:
		return lua.LNumber(v.Value().(uint32))
	case ole.VT_I8:
		return lua.LNumber(v.Value().(int64))
	case ole.VT_UI8:
		return lua.LNumber(v.Value().(uint64))
	case ole.VT_INT:
		return lua.LNumber(v.Value().(int))
	case ole.VT_UINT:
		return lua.LNumber(v.Value().(uint))
	case ole.VT_INT_PTR:
		return lua.LNumber(v.Value().(uintptr))
	case ole.VT_UINT_PTR:
		return lua.LNumber(v.Value().(uintptr))
	case ole.VT_R4:
		return lua.LNumber(v.Value().(float32))
	case ole.VT_R8:
		return lua.LNumber(v.Value().(float64))
	case ole.VT_BSTR:
		return lua.LString(v.ToString())
	case ole.VT_DATE:
		if date, ok := v.Value().(time.Time); ok {
			t := L.NewTable()
			L.SetField(t, "year", lua.LNumber(date.Year()))
			L.SetField(t, "month", lua.LNumber(int(date.Month())))
			L.SetField(t, "day", lua.LNumber(date.Day()))
			L.SetField(t, "hour", lua.LNumber(date.Hour()))
			L.SetField(t, "min", lua.LNumber(date.Minute()))
			L.SetField(t, "sec", lua.LNumber(date.Second()))
			return t
		} else if floatValue, ok := v.Value().(float64); ok {
			return lua.LNumber(floatValue)
		} else {
			panic("can not convert ole.VT_DATE")
		}
	case ole.VT_UNKNOWN:
		panic("can not convert ole.VT_UNKNOW")
	case ole.VT_DISPATCH:
		return capsuleT{v.ToIDispatch()}.ToLValue(L)
	case ole.VT_BOOL:
		if v.Value().(bool) {
			return lua.LTrue
		} else {
			return lua.LFalse
		}
	default:
		return lua.LNil
	}
}
<|endoftext|>