<commit_before>package main

import (
    "gopkg.in/yaml.v2"
    "io/ioutil"
    "fmt"
    "encoding/hex"
    "encoding/binary"
    "github.com/fuzxxl/nfc/2.0/nfc"
    "github.com/fuzxxl/freefare/0.3/freefare"
    "./keydiversification"
)

// TODO: move to a separate helper module
func string_to_aeskey(keydata_str string) (*freefare.DESFireKey, error) {
    keydata := new([16]byte)
    to_keydata, err := hex.DecodeString(keydata_str)
    if err != nil {
        key := freefare.NewDESFireAESKey(*keydata, 0)
        return key, err
    }
    copy(keydata[0:], to_keydata)
    key := freefare.NewDESFireAESKey(*keydata, 0)
    return key,nil
}

func bytes_to_aeskey(source []byte) (*freefare.DESFireKey) {
    keydata := new([16]byte)
    copy(keydata[0:], source)
    key := freefare.NewDESFireAESKey(*keydata, 0)
    return key
}

func string_to_byte(source string) (byte, error) {
    bytearray, err := hex.DecodeString(source)
    if err != nil {
        return 0x0, err
    }
    return bytearray[0], nil
}

func main() {
    keys_data, err := ioutil.ReadFile("keys.yaml")
    if err != nil {
        panic(err)
    }

    keymap := make(map[interface{}]interface{});
    err = yaml.Unmarshal([]byte(keys_data), &keymap);
    if err != nil {
        panic(err)
    }

    apps_data, err := ioutil.ReadFile("apps.yaml")
    if err != nil {
        panic(err)
    }

    appmap := make(map[interface{}]interface{});
    err = yaml.Unmarshal([]byte(apps_data), &appmap);
    if err != nil {
        panic(err)
    }

    // Application-id from config
    aidbytes, err := hex.DecodeString(appmap["hacklab_acl"].(map[interface{}]interface{})["aid"].(string))
    if err != nil {
        panic(err)
    }
    aidint, n := binary.Uvarint(aidbytes)
    if n <= 0 {
        panic(fmt.Sprintf("binary.Uvarint returned %d", n))
    }
    aid := freefare.NewDESFireAid(uint32(aidint))
    //fmt.Println(aid)
    // Needed for diversification
    sysid, err := hex.DecodeString(appmap["hacklab_acl"].(map[interface{}]interface{})["sysid"].(string))
    if err != nil {
        panic(err)
    }

    // Key id numbers from config
    uid_read_key_id, err := string_to_byte(appmap["hacklab_acl"].(map[interface{}]interface{})["uid_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    acl_read_key_id, err := string_to_byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    acl_write_key_id, err := string_to_byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_write_key_id"].(string))
    if err != nil {
        panic(err)
    }
    prov_key_id, err := string_to_byte(appmap["hacklab_acl"].(map[interface{}]interface{})["provisioning_key_id"].(string))
    if err != nil {
        panic(err)
    }

    // Defaul (null) key
    nullkeydata := new([8]byte)
    defaultkey := freefare.NewDESFireDESKey(*nullkeydata)

    // New card master key
    new_master_key, err := string_to_aeskey(keymap["card_master"].(string))
    if err != nil {
        panic(err)
    }
    //fmt.Println(new_master_key)

    // The static app key to read UID
    uid_read_key, err := string_to_aeskey(keymap["uid_read_key"].(string))
    if err != nil {
        panic(err)
    }
    //fmt.Println(uid_read_key)

    // Bases for the diversified keys
    prov_key_base, err := hex.DecodeString(keymap["prov_master"].(string))
    if err != nil {
        panic(err)
    }
    acl_read_base, err := hex.DecodeString(keymap["acl_read_key"].(string))
    if err != nil {
        panic(err)
    }
    acl_write_base, err := hex.DecodeString(keymap["acl_write_key"].(string))
    if err != nil {
        panic(err)
    }

    // Open device and get tags list
    d, err := nfc.Open("");
    if err != nil {
        panic(err)
    }

    tags, err := freefare.GetTags(d);
    if err != nil {
        panic(err)
    }

    // Initialize each tag with our app
    for i := 0; i < len(tags); i++ {
        tag := tags[i]
        fmt.Println(tag.String(), tag.UID())

        // Skip non desfire tags
        if (tag.Type() != freefare.DESFire) {
            fmt.Println("Skipped");
            continue
        }

        desfiretag := tag.(freefare.DESFireTag)

        // Connect to this tag
        fmt.Println("Connecting");
        error := desfiretag.Connect()
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Authenticating");
        error = desfiretag.Authenticate(0,*defaultkey)
        if error != nil {
            fmt.Println("Failed, trying agin with new key")
            error = desfiretag.Authenticate(0,*new_master_key)
            if error != nil {
                panic(error)
            }
            fmt.Println("Changing key back to default")
            error = desfiretag.ChangeKey(0,  *defaultkey, *new_master_key);
            if error != nil {
                panic(error)
            }
            fmt.Println("Re-auth with default key")
            error = desfiretag.Authenticate(0,*defaultkey)
            if error != nil {
                panic(error)
            }
            fmt.Println("Formatting (to get a clean state)")
            error = desfiretag.FormatPICC()
            if error != nil {
                panic(error)
            }
            return
        }
        fmt.Println("Done");

        // Get card real UID
        realuid_str, error := desfiretag.CardUID()
        if error != nil {
            panic(error)
        }
        realuid, error := hex.DecodeString(realuid_str);
        if error != nil {
            panic(error)
        }

        // Calculate the diversified keys
        prov_key_bytes, err := keydiversification.AES128(prov_key_base, aidbytes, realuid, sysid)
        if err != nil {
            panic(err)
        }
        prov_key := bytes_to_aeskey(prov_key_bytes)
        acl_read_bytes, err := keydiversification.AES128(acl_read_base, aidbytes, realuid, sysid)
        if err != nil {
            panic(err)
        }
        acl_read_key := bytes_to_aeskey(acl_read_bytes)
        acl_write_bytes, err := keydiversification.AES128(acl_write_base, aidbytes, realuid, sysid)
        if err != nil {
            panic(err)
        }
        acl_write_key := bytes_to_aeskey(acl_write_bytes)

        fmt.Println("Changing default master key");
        error = desfiretag.ChangeKey(0, *new_master_key, *defaultkey);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        /**
         * This is not needed for creating the application and does not help when changing application keys
        fmt.Println("Re-auth with new key")
        error = desfiretag.Authenticate(0,*new_master_key)
        if error != nil {
            panic(error)
        }
         */

        fmt.Println("Creating application");
        // TODO:Figure out what the settings byte (now hardcoded to 0xFF as it was in libfreefare example code) actually does
        error = desfiretag.CreateApplication(aid, 0xFF, 6 | freefare.CryptoAES);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Selecting application");
        // TODO:Figure out what the settings byte (now hardcoded to 0xFF as it was in libfreefare exampkle code) actually does
        error = desfiretag.SelectApplication(aid);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        /**
         * Does not work
        fmt.Println("Re-auth with new master key")
        error = desfiretag.Authenticate(0,*new_master_key)
        if error != nil {
            panic(error)
        }
         */

        /**
         * Also does not work
        fmt.Println("Re-auth with default key")
        error = desfiretag.Authenticate(0,*defaultkey)
        if error != nil {
            panic(error)
        }
         */

        fmt.Println("Changing provisioning key");
        error = desfiretag.ChangeKey(prov_key_id, *prov_key, *defaultkey);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Re-auth with new provisioning key")
        error = desfiretag.Authenticate(prov_key_id,*prov_key)
        if error != nil {
            panic(error)
        }

        fmt.Println("Changing static UID reading key");
        error = desfiretag.ChangeKey(uid_read_key_id, *uid_read_key, *defaultkey);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Changing ACL reading key");
        error = desfiretag.ChangeKey(acl_read_key_id, *acl_read_key, *defaultkey);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Changing ACL writing key");
        error = desfiretag.ChangeKey(acl_write_key_id, *acl_write_key, *defaultkey);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Creating ACL data file");
        error = desfiretag.CreateDataFile(0, freefare.Enciphered, freefare.MakeDESFireAccessRights(acl_read_key_id, acl_write_key_id, prov_key_id, prov_key_id), 8, false)
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        /**
         * TODO: check later if this works
        fmt.Println("Creating ACL backup file");
        error = desfiretag.CreateDataFile(0, freefare.Enciphered, freefare.MakeDESFireAccessRights(acl_read_key_id, acl_write_key_id, prov_key_id, prov_key_id), 8, true)
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");
         */

        /**
         * Only needed when working with backup files
        // Not sure if this is actually needed
        fmt.Println("Committing");
        error = desfiretag.CommitTransaction()
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");
         */

        fmt.Println("Disconnecting");
        error = desfiretag.Disconnect()
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");
    }

}<commit_msg>null AES key auth before changing application key 0 (aka provisioning key) worked, but getting problems when trying to change the other keys<commit_after>package main

import (
    "gopkg.in/yaml.v2"
    "io/ioutil"
    "fmt"
    "encoding/hex"
    "encoding/binary"
    "github.com/fuzxxl/nfc/2.0/nfc"
    "github.com/fuzxxl/freefare/0.3/freefare"
    "./keydiversification"
)

// TODO: move to a separate helper module
func string_to_aeskey(keydata_str string) (*freefare.DESFireKey, error) {
    keydata := new([16]byte)
    to_keydata, err := hex.DecodeString(keydata_str)
    if err != nil {
        key := freefare.NewDESFireAESKey(*keydata, 0)
        return key, err
    }
    copy(keydata[0:], to_keydata)
    key := freefare.NewDESFireAESKey(*keydata, 0)
    return key,nil
}

func bytes_to_aeskey(source []byte) (*freefare.DESFireKey) {
    keydata := new([16]byte)
    copy(keydata[0:], source)
    key := freefare.NewDESFireAESKey(*keydata, 0)
    return key
}

func string_to_byte(source string) (byte, error) {
    bytearray, err := hex.DecodeString(source)
    if err != nil {
        return 0x0, err
    }
    return bytearray[0], nil
}

func main() {
    keys_data, err := ioutil.ReadFile("keys.yaml")
    if err != nil {
        panic(err)
    }

    keymap := make(map[interface{}]interface{});
    err = yaml.Unmarshal([]byte(keys_data), &keymap);
    if err != nil {
        panic(err)
    }

    apps_data, err := ioutil.ReadFile("apps.yaml")
    if err != nil {
        panic(err)
    }

    appmap := make(map[interface{}]interface{});
    err = yaml.Unmarshal([]byte(apps_data), &appmap);
    if err != nil {
        panic(err)
    }

    // Application-id from config
    aidbytes, err := hex.DecodeString(appmap["hacklab_acl"].(map[interface{}]interface{})["aid"].(string))
    if err != nil {
        panic(err)
    }
    aidint, n := binary.Uvarint(aidbytes)
    if n <= 0 {
        panic(fmt.Sprintf("binary.Uvarint returned %d", n))
    }
    aid := freefare.NewDESFireAid(uint32(aidint))
    //fmt.Println(aid)
    // Needed for diversification
    sysid, err := hex.DecodeString(appmap["hacklab_acl"].(map[interface{}]interface{})["sysid"].(string))
    if err != nil {
        panic(err)
    }

    // Key id numbers from config
    uid_read_key_id, err := string_to_byte(appmap["hacklab_acl"].(map[interface{}]interface{})["uid_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    acl_read_key_id, err := string_to_byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    acl_write_key_id, err := string_to_byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_write_key_id"].(string))
    if err != nil {
        panic(err)
    }
    prov_key_id, err := string_to_byte(appmap["hacklab_acl"].(map[interface{}]interface{})["provisioning_key_id"].(string))
    if err != nil {
        panic(err)
    }

    // Defaul (null) key
    nullkeydata := new([8]byte)
    defaultkey := freefare.NewDESFireDESKey(*nullkeydata)

    nullkeydata16 := new([16]byte)
    defaultkey_aes := freefare.NewDESFireAESKey(*nullkeydata16, 0)

    // New card master key
    new_master_key, err := string_to_aeskey(keymap["card_master"].(string))
    if err != nil {
        panic(err)
    }
    //fmt.Println(new_master_key)

    // The static app key to read UID
    uid_read_key, err := string_to_aeskey(keymap["uid_read_key"].(string))
    if err != nil {
        panic(err)
    }
    //fmt.Println(uid_read_key)

    // Bases for the diversified keys
    prov_key_base, err := hex.DecodeString(keymap["prov_master"].(string))
    if err != nil {
        panic(err)
    }
    acl_read_base, err := hex.DecodeString(keymap["acl_read_key"].(string))
    if err != nil {
        panic(err)
    }
    acl_write_base, err := hex.DecodeString(keymap["acl_write_key"].(string))
    if err != nil {
        panic(err)
    }

    // Open device and get tags list
    d, err := nfc.Open("");
    if err != nil {
        panic(err)
    }

    tags, err := freefare.GetTags(d);
    if err != nil {
        panic(err)
    }

    // Initialize each tag with our app
    for i := 0; i < len(tags); i++ {
        tag := tags[i]
        fmt.Println(tag.String(), tag.UID())

        // Skip non desfire tags
        if (tag.Type() != freefare.DESFire) {
            fmt.Println("Skipped");
            continue
        }

        desfiretag := tag.(freefare.DESFireTag)

        // Connect to this tag
        fmt.Println("Connecting");
        error := desfiretag.Connect()
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Authenticating");
        error = desfiretag.Authenticate(0,*defaultkey)
        if error != nil {
            fmt.Println("Failed, trying agin with new key")
            error = desfiretag.Authenticate(0,*new_master_key)
            if error != nil {
                panic(error)
            }
            fmt.Println("Changing key back to default")
            error = desfiretag.ChangeKey(0,  *defaultkey, *new_master_key);
            if error != nil {
                panic(error)
            }
            fmt.Println("Re-auth with default key")
            error = desfiretag.Authenticate(0,*defaultkey)
            if error != nil {
                panic(error)
            }
            fmt.Println("Formatting (to get a clean state)")
            error = desfiretag.FormatPICC()
            if error != nil {
                panic(error)
            }
            return
        }
        fmt.Println("Done");

        // Get card real UID
        realuid_str, error := desfiretag.CardUID()
        if error != nil {
            panic(error)
        }
        realuid, error := hex.DecodeString(realuid_str);
        if error != nil {
            panic(error)
        }

        // Calculate the diversified keys
        prov_key_bytes, err := keydiversification.AES128(prov_key_base, aidbytes, realuid, sysid)
        if err != nil {
            panic(err)
        }
        prov_key := bytes_to_aeskey(prov_key_bytes)
        acl_read_bytes, err := keydiversification.AES128(acl_read_base, aidbytes, realuid, sysid)
        if err != nil {
            panic(err)
        }
        acl_read_key := bytes_to_aeskey(acl_read_bytes)
        acl_write_bytes, err := keydiversification.AES128(acl_write_base, aidbytes, realuid, sysid)
        if err != nil {
            panic(err)
        }
        acl_write_key := bytes_to_aeskey(acl_write_bytes)

        fmt.Println("Changing default master key");
        error = desfiretag.ChangeKey(0, *new_master_key, *defaultkey);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        /**
         * This is not needed for creating the application and does not help when changing application keys
        fmt.Println("Re-auth with new key")
        error = desfiretag.Authenticate(0,*new_master_key)
        if error != nil {
            panic(error)
        }
         */

        fmt.Println("Creating application");
        // TODO:Figure out what the settings byte (now hardcoded to 0xFF as it was in libfreefare example code) actually does
        error = desfiretag.CreateApplication(aid, 0xFF, 6 | freefare.CryptoAES);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Selecting application");
        // TODO:Figure out what the settings byte (now hardcoded to 0xFF as it was in libfreefare exampkle code) actually does
        error = desfiretag.SelectApplication(aid);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        /**
         * Does not work
        fmt.Println("Re-auth with new master key")
        error = desfiretag.Authenticate(0,*new_master_key)
        if error != nil {
            panic(error)
        }
         */

        /**
         * Also does not work
        fmt.Println("Re-auth with default key")
        error = desfiretag.Authenticate(0,*defaultkey)
        if error != nil {
            panic(error)
        }
         */

        fmt.Println("Re-auth with null AES key")
        error = desfiretag.Authenticate(prov_key_id,*defaultkey_aes)
        if error != nil {
            panic(error)
        }
        fmt.Println("Changing provisioning key");
        error = desfiretag.ChangeKey(prov_key_id, *prov_key, *defaultkey_aes);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Re-auth with null AES key")
        error = desfiretag.Authenticate(uid_read_key_id,*defaultkey_aes)
        if error != nil {
            panic(error)
        }
        fmt.Println("Changing static UID reading key");
        error = desfiretag.ChangeKey(uid_read_key_id, *uid_read_key, *defaultkey_aes);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Re-auth with null AES key")
        error = desfiretag.Authenticate(acl_read_key_id,*defaultkey_aes)
        if error != nil {
            panic(error)
        }
        fmt.Println("Changing ACL reading key");
        error = desfiretag.ChangeKey(acl_read_key_id, *acl_read_key, *defaultkey_aes);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Re-auth with null AES key")
        error = desfiretag.Authenticate(acl_write_key_id,*defaultkey_aes)
        if error != nil {
            panic(error)
        }
        fmt.Println("Changing ACL writing key");
        error = desfiretag.ChangeKey(acl_write_key_id, *acl_write_key, *defaultkey_aes);
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        fmt.Println("Re-auth with provisioning key")
        error = desfiretag.Authenticate(prov_key_id,*prov_key)
        if error != nil {
            panic(error)
        }

        fmt.Println("Creating ACL data file");
        error = desfiretag.CreateDataFile(0, freefare.Enciphered, freefare.MakeDESFireAccessRights(acl_read_key_id, acl_write_key_id, prov_key_id, prov_key_id), 8, false)
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");

        /**
         * TODO: check later if this works
        fmt.Println("Creating ACL backup file");
        error = desfiretag.CreateDataFile(0, freefare.Enciphered, freefare.MakeDESFireAccessRights(acl_read_key_id, acl_write_key_id, prov_key_id, prov_key_id), 8, true)
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");
         */

        /**
         * Only needed when working with backup files
        // Not sure if this is actually needed
        fmt.Println("Committing");
        error = desfiretag.CommitTransaction()
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");
         */

        fmt.Println("Disconnecting");
        error = desfiretag.Disconnect()
        if error != nil {
            panic(error)
        }
        fmt.Println("Done");
    }

}<|endoftext|>


<commit_before>// Copyright 2013 The Pythia Authors.
// This file is part of Pythia.
//
// Pythia is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, version 3 of the License.
//
// Pythia is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with Pythia.  If not, see <http://www.gnu.org/licenses/>.

package pythia

import (
	"encoding/json"
	"errors"
	"io"
	"log"
	"net"
	"time"
)

// Error to return if the connection was closed
var closedError = errors.New("Connection closed")

// MessageResult is an auxiliary structure for passing messages to the writer
// goroutine.
type messageResult struct {
	Msg    Message      // what to send
	Result chan<- error // where to write the result of the operation
}

// Conn is a wrapper over net.Conn, reading and writing Messages.
type Conn struct {
	// The underlying connection.
	conn net.Conn

	// Incoming messages channel.
	input chan Message

	// Outgoing messages channel.
	output chan messageResult

	// Channel to ask reader and writer goroutines to quit.
	quit chan bool

	// Flag to ignore errors after closing the connection.
	closed bool
}

// WrapConn wraps a stream network connection into a Message-oriented
// connection. The raw conn connection shall not be used by the user anymore.
func WrapConn(conn net.Conn) *Conn {
	c := new(Conn)
	c.conn = conn
	c.input = make(chan Message)
	c.output = make(chan messageResult)
	c.quit = make(chan bool, 1)
	go c.reader()
	go c.writer()
	return c
}

// The reader goroutine parses the Messages and put them in the input channel.
// Keep-alive messages are discarded.
func (c *Conn) reader() {
	defer close(c.input)
	dec := json.NewDecoder(c.conn)
	for {
		// Note: As a keep-alive message is sent only when no message has been
		// transmitted during the previous interval, it is possible that no
		// message gets transmitted for nearly 2 intervals. To take this into
		// account (plus some safety margin), we set the read timeout to
		// 3 intervals.
		c.conn.SetReadDeadline(time.Now().Add(3 * KeepAliveInterval))
		var msg Message
		err := dec.Decode(&msg)
		if c.closed {
			return
		} else if err == io.EOF {
			log.Println("Connection closed on remote side.")
			c.Close()
			return
		} else if neterr, ok := err.(net.Error); ok && neterr.Timeout() {
			log.Println("Connection timed out.")
			c.Close()
			return
		} else if err != nil {
			log.Print(err)
		} else if msg.Message != KeepAliveMsg {
			c.input <- msg
		}
	}
}

// The writer goroutine sends Messages and keep-alives
func (c *Conn) writer() {
	keepAliveTicker := time.NewTicker(KeepAliveInterval)
	defer keepAliveTicker.Stop()
	sendKeepAlive := true
	enc := json.NewEncoder(c.conn)
	for {
		select {
		case mr := <-c.output:
			msg, result := mr.Msg, mr.Result
			result <- enc.Encode(msg)
			sendKeepAlive = false
		case <-keepAliveTicker.C:
			if sendKeepAlive {
				err := enc.Encode(Message{Message: KeepAliveMsg})
				if err != nil {
					log.Println("Error sending keep-alive message:", err)
				}
			}
			sendKeepAlive = true
		case <-c.quit:
			c.sendQuit()
			return
		}
	}
}

// Dial connects to the address addr and returns a Message-oriented connection.
func Dial(addr net.Addr) (*Conn, error) {
	conn, err := net.Dial(addr.Network(), addr.String())
	if err != nil {
		return nil, err
	}
	return WrapConn(conn), nil
}

// DialRetry is equivalent to Dial, but it keeps on retrying (with exponential
// back-off) until the connection is established.
func DialRetry(addr net.Addr) *Conn {
	interval := InitialRetryInterval
	for {
		conn, err := Dial(addr)
		if err == nil {
			return conn
		}
		log.Print("Connection failed (", err, "), retrying in ", interval)
		time.Sleep(interval)
		interval *= 2
		if interval > MaxRetryInterval {
			interval = MaxRetryInterval
		}
	}
}

// Receive returns the channel from which incoming messages can be retrieved.
// The channel is closed when the connection is closed.
func (c *Conn) Receive() <-chan Message {
	return c.input
}

// Send sends a message through the connection.
func (c *Conn) Send(msg Message) error {
	if c.closed {
		return closedError
	}
	result := make(chan error)
	select {
	case c.output <- messageResult{Msg: msg, Result: result}:
	case <-c.quit:
		c.sendQuit()
		return closedError
	}
	select {
	case err := <-result:
		return err
	case <-c.quit:
		c.sendQuit()
		return closedError
	}
}

// sendQuit signals the quit channel, but does not block. This requires the
// quit channel to be buffered.
func (c *Conn) sendQuit() {
	select {
	case c.quit <- true:
	default:
	}
}

// Close closes the connection. The receive channel will also be closed.
// Further sends will cause errors.
func (c *Conn) Close() error {
	c.sendQuit()
	c.closed = true
	return c.conn.Close()
}

// vim:set sw=4 ts=4 noet:
<commit_msg>Go: Remove spurious error message<commit_after>// Copyright 2013-2020 The Pythia Authors.
// This file is part of Pythia.
//
// Pythia is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, version 3 of the License.
//
// Pythia is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with Pythia.  If not, see <http://www.gnu.org/licenses/>.

package pythia

import (
	"encoding/json"
	"errors"
	"io"
	"log"
	"net"
	"time"
)

// Error to return if the connection was closed
var closedError = errors.New("Connection closed")

// MessageResult is an auxiliary structure for passing messages to the writer
// goroutine.
type messageResult struct {
	Msg    Message      // what to send
	Result chan<- error // where to write the result of the operation
}

// Conn is a wrapper over net.Conn, reading and writing Messages.
type Conn struct {
	// The underlying connection.
	conn net.Conn

	// Incoming messages channel.
	input chan Message

	// Outgoing messages channel.
	output chan messageResult

	// Channel to ask reader and writer goroutines to quit.
	quit chan bool

	// Flag to ignore errors after closing the connection.
	closed bool
}

// WrapConn wraps a stream network connection into a Message-oriented
// connection. The raw conn connection shall not be used by the user anymore.
func WrapConn(conn net.Conn) *Conn {
	c := new(Conn)
	c.conn = conn
	c.input = make(chan Message)
	c.output = make(chan messageResult)
	c.quit = make(chan bool, 1)
	go c.reader()
	go c.writer()
	return c
}

// The reader goroutine parses the Messages and put them in the input channel.
// Keep-alive messages are discarded.
func (c *Conn) reader() {
	defer close(c.input)
	dec := json.NewDecoder(c.conn)
	for {
		// Note: As a keep-alive message is sent only when no message has been
		// transmitted during the previous interval, it is possible that no
		// message gets transmitted for nearly 2 intervals. To take this into
		// account (plus some safety margin), we set the read timeout to
		// 3 intervals.
		c.conn.SetReadDeadline(time.Now().Add(3 * KeepAliveInterval))
		var msg Message
		err := dec.Decode(&msg)
		if c.closed {
			return
		} else if err == io.EOF {
			log.Println("Connection closed on remote side.")
			c.Close()
			return
		} else if neterr, ok := err.(net.Error); ok && neterr.Timeout() {
			log.Println("Connection timed out.")
			c.Close()
			return
		} else if err != nil {
			log.Print(err)
			return
		} else if msg.Message != KeepAliveMsg {
			c.input <- msg
		}
	}
}

// The writer goroutine sends Messages and keep-alives
func (c *Conn) writer() {
	keepAliveTicker := time.NewTicker(KeepAliveInterval)
	defer keepAliveTicker.Stop()
	sendKeepAlive := true
	enc := json.NewEncoder(c.conn)
	for {
		select {
		case mr := <-c.output:
			msg, result := mr.Msg, mr.Result
			result <- enc.Encode(msg)
			sendKeepAlive = false
		case <-keepAliveTicker.C:
			if sendKeepAlive {
				err := enc.Encode(Message{Message: KeepAliveMsg})
				if err != nil {
					log.Println("Error sending keep-alive message:", err)
				}
			}
			sendKeepAlive = true
		case <-c.quit:
			c.sendQuit()
			return
		}
	}
}

// Dial connects to the address addr and returns a Message-oriented connection.
func Dial(addr net.Addr) (*Conn, error) {
	conn, err := net.Dial(addr.Network(), addr.String())
	if err != nil {
		return nil, err
	}
	return WrapConn(conn), nil
}

// DialRetry is equivalent to Dial, but it keeps on retrying (with exponential
// back-off) until the connection is established.
func DialRetry(addr net.Addr) *Conn {
	interval := InitialRetryInterval
	for {
		conn, err := Dial(addr)
		if err == nil {
			return conn
		}
		log.Print("Connection failed (", err, "), retrying in ", interval)
		time.Sleep(interval)
		interval *= 2
		if interval > MaxRetryInterval {
			interval = MaxRetryInterval
		}
	}
}

// Receive returns the channel from which incoming messages can be retrieved.
// The channel is closed when the connection is closed.
func (c *Conn) Receive() <-chan Message {
	return c.input
}

// Send sends a message through the connection.
func (c *Conn) Send(msg Message) error {
	if c.closed {
		return closedError
	}
	result := make(chan error)
	select {
	case c.output <- messageResult{Msg: msg, Result: result}:
	case <-c.quit:
		c.sendQuit()
		return closedError
	}
	select {
	case err := <-result:
		return err
	case <-c.quit:
		c.sendQuit()
		return closedError
	}
}

// sendQuit signals the quit channel, but does not block. This requires the
// quit channel to be buffered.
func (c *Conn) sendQuit() {
	select {
	case c.quit <- true:
	default:
	}
}

// Close closes the connection. The receive channel will also be closed.
// Further sends will cause errors.
func (c *Conn) Close() error {
	c.sendQuit()
	c.closed = true
	return c.conn.Close()
}

// vim:set sw=4 ts=4 noet:
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package goapp

import (
	"bytes"
	"code.google.com/p/go.net/html"
	"code.google.com/p/go.net/html/atom"
	"errors"
	"io"
)

var ErrNoRssLink = errors.New("No rss link found")

func Autodiscover(b []byte) (string, error) {
	r := bytes.NewReader(b)
	z := html.NewTokenizer(r)
	inHtml := false
	inHead := false
	for {
		if z.Next() == html.ErrorToken {
			if err := z.Err(); err == io.EOF {
				break
			} else {
				return "done", ErrNoRssLink
			}
		}
		t := z.Token()
		switch t.DataAtom {
		case atom.Html:
			inHtml = !inHtml
		case atom.Head:
			inHead = !inHead
		case atom.Link:
			if inHead && inHtml && (t.Type == html.StartTagToken || t.Type == html.SelfClosingTagToken) {
				attrs := make(map[string]string)
				for _, a := range t.Attr {
					attrs[a.Key] = a.Val
				}
				if attrs["rel"] == "alternate" && attrs["type"] == "application/rss+xml" && attrs["href"] != "" {
					return attrs["href"], nil
				}
			}
		}
	}

	return "end", ErrNoRssLink
}
<commit_msg>Remove debugs<commit_after>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package goapp

import (
	"bytes"
	"code.google.com/p/go.net/html"
	"code.google.com/p/go.net/html/atom"
	"errors"
	"io"
)

var ErrNoRssLink = errors.New("No rss link found")

func Autodiscover(b []byte) (string, error) {
	r := bytes.NewReader(b)
	z := html.NewTokenizer(r)
	inHtml := false
	inHead := false
	for {
		if z.Next() == html.ErrorToken {
			if err := z.Err(); err == io.EOF {
				break
			} else {
				return "", ErrNoRssLink
			}
		}
		t := z.Token()
		switch t.DataAtom {
		case atom.Html:
			inHtml = !inHtml
		case atom.Head:
			inHead = !inHead
		case atom.Link:
			if inHead && inHtml && (t.Type == html.StartTagToken || t.Type == html.SelfClosingTagToken) {
				attrs := make(map[string]string)
				for _, a := range t.Attr {
					attrs[a.Key] = a.Val
				}
				if attrs["rel"] == "alternate" && attrs["type"] == "application/rss+xml" && attrs["href"] != "" {
					return attrs["href"], nil
				}
			}
		}
	}

	return "", ErrNoRssLink
}
<|endoftext|>


<commit_before>package gocbcore

const (
	GoCbVersionStr = "v1.0.8"
)

type CommandMagic uint8

const (
	ReqMagic = CommandMagic(0x80)
	ResMagic = CommandMagic(0x81)
)

// CommandCode for memcached packets.
type CommandCode uint8

const (
	CmdGet                  = CommandCode(0x00)
	CmdSet                  = CommandCode(0x01)
	CmdAdd                  = CommandCode(0x02)
	CmdReplace              = CommandCode(0x03)
	CmdDelete               = CommandCode(0x04)
	CmdIncrement            = CommandCode(0x05)
	CmdDecrement            = CommandCode(0x06)
	CmdAppend               = CommandCode(0x0e)
	CmdPrepend              = CommandCode(0x0f)
	CmdStat                 = CommandCode(0x10)
	CmdTouch                = CommandCode(0x1c)
	CmdGAT                  = CommandCode(0x1d)
	CmdHello                = CommandCode(0x1f)
	CmdSASLListMechs        = CommandCode(0x20)
	CmdSASLAuth             = CommandCode(0x21)
	CmdSASLStep             = CommandCode(0x22)
	CmdGetAllVBSeqnos       = CommandCode(0x48)
	CmdDcpOpenConnection    = CommandCode(0x50)
	CmdDcpAddStream         = CommandCode(0x51)
	CmdDcpCloseStream       = CommandCode(0x52)
	CmdDcpStreamReq         = CommandCode(0x53)
	CmdDcpGetFailoverLog    = CommandCode(0x54)
	CmdDcpStreamEnd         = CommandCode(0x55)
	CmdDcpSnapshotMarker    = CommandCode(0x56)
	CmdDcpMutation          = CommandCode(0x57)
	CmdDcpDeletion          = CommandCode(0x58)
	CmdDcpExpiration        = CommandCode(0x59)
	CmdDcpFlush             = CommandCode(0x5a)
	CmdDcpSetVbucketState   = CommandCode(0x5b)
	CmdDcpNoop              = CommandCode(0x5c)
	CmdDcpBufferAck         = CommandCode(0x5d)
	CmdDcpControl           = CommandCode(0x5e)
	CmdGetReplica           = CommandCode(0x83)
	CmdSelectBucket         = CommandCode(0x89)
	CmdObserveSeqNo         = CommandCode(0x91)
	CmdObserve              = CommandCode(0x92)
	CmdGetLocked            = CommandCode(0x94)
	CmdUnlockKey            = CommandCode(0x95)
	CmdSetMeta              = CommandCode(0xa2)
	CmdDetMeta              = CommandCode(0xa8)
	CmdGetClusterConfig     = CommandCode(0xb5)
	CmdGetRandom            = CommandCode(0xb6)
	CmdSubDocGet            = CommandCode(0xc5)
	CmdSubDocExists         = CommandCode(0xc6)
	CmdSubDocDictAdd        = CommandCode(0xc7)
	CmdSubDocDictSet        = CommandCode(0xc8)
	CmdSubDocDelete         = CommandCode(0xc9)
	CmdSubDocReplace        = CommandCode(0xca)
	CmdSubDocArrayPushLast  = CommandCode(0xcb)
	CmdSubDocArrayPushFirst = CommandCode(0xcc)
	CmdSubDocArrayInsert    = CommandCode(0xcd)
	CmdSubDocArrayAddUnique = CommandCode(0xce)
	CmdSubDocCounter        = CommandCode(0xcf)
	CmdSubDocMultiLookup    = CommandCode(0xd0)
	CmdSubDocMultiMutation  = CommandCode(0xd1)
)

type SubDocFlag uint16

const (
	SubDocFlagMkDirP = SubDocFlag(0x01)
)

type SubDocOpType uint8

const (
	SubDocOpGet            = SubDocOpType(CmdSubDocGet)
	SubDocOpExists         = SubDocOpType(CmdSubDocExists)
	SubDocOpDictAdd        = SubDocOpType(CmdSubDocDictAdd)
	SubDocOpDictSet        = SubDocOpType(CmdSubDocDictSet)
	SubDocOpDelete         = SubDocOpType(CmdSubDocDelete)
	SubDocOpReplace        = SubDocOpType(CmdSubDocReplace)
	SubDocOpArrayPushLast  = SubDocOpType(CmdSubDocArrayPushLast)
	SubDocOpArrayPushFirst = SubDocOpType(CmdSubDocArrayPushFirst)
	SubDocOpArrayInsert    = SubDocOpType(CmdSubDocArrayInsert)
	SubDocOpArrayAddUnique = SubDocOpType(CmdSubDocArrayAddUnique)
	SubDocOpCounter        = SubDocOpType(CmdSubDocCounter)
)

type HelloFeature uint16

const (
	FeatureDatatype = HelloFeature(0x01)
	FeatureSeqNo    = HelloFeature(0x04)
)

// Status field for memcached response.
type StatusCode uint16

const (
	StatusSuccess            = StatusCode(0x00)
	StatusKeyNotFound        = StatusCode(0x01)
	StatusKeyExists          = StatusCode(0x02)
	StatusTooBig             = StatusCode(0x03)
	StatusInvalidArgs        = StatusCode(0x04)
	StatusNotStored          = StatusCode(0x05)
	StatusBadDelta           = StatusCode(0x06)
	StatusNotMyVBucket       = StatusCode(0x07)
	StatusNoBucket           = StatusCode(0x08)
	StatusAuthStale          = StatusCode(0x1f)
	StatusAuthError          = StatusCode(0x20)
	StatusAuthContinue       = StatusCode(0x21)
	StatusRangeError         = StatusCode(0x22)
	StatusRollback           = StatusCode(0x23)
	StatusAccessError        = StatusCode(0x24)
	StatusNotInitialized     = StatusCode(0x25)
	StatusUnknownCommand     = StatusCode(0x81)
	StatusOutOfMemory        = StatusCode(0x82)
	StatusNotSupported       = StatusCode(0x83)
	StatusInternalError      = StatusCode(0x84)
	StatusBusy               = StatusCode(0x85)
	StatusTmpFail            = StatusCode(0x86)
	StatusSubDocPathNotFound = StatusCode(0xc0)
	StatusSubDocPathMismatch = StatusCode(0xc1)
	StatusSubDocPathInvalid  = StatusCode(0xc2)
	StatusSubDocPathTooBig   = StatusCode(0xc3)
	StatusSubDocDocTooDeep   = StatusCode(0xc4)
	StatusSubDocCantInsert   = StatusCode(0xc5)
	StatusSubDocNotJson      = StatusCode(0xc6)
	StatusSubDocBadRange     = StatusCode(0xc7)
	StatusSubDocBadDelta     = StatusCode(0xc8)
	StatusSubDocPathExists   = StatusCode(0xc9)
	StatusSubDocValueTooDeep = StatusCode(0xca)
	StatusSubDocBadCombo     = StatusCode(0xcb)
	StatusSubDocBadMulti     = StatusCode(0xcc)
)

type KeyState uint8

const (
	KeyStateNotPersisted = KeyState(0x00)
	KeyStatePersisted    = KeyState(0x01)
	KeyStateNotFound     = KeyState(0x80)
	KeyStateDeleted      = KeyState(0x81)
)

type StreamEndStatus uint32

const (
	StreamEndOK           = StreamEndStatus(0x00)
	StreamEndClosed       = StreamEndStatus(0x01)
	StreamEndStateChanged = StreamEndStatus(0x02)
	StreamEndDisconnected = StreamEndStatus(0x03)
	StreamEndTooSlow      = StreamEndStatus(0x04)
)

type BucketType int

const (
	BktTypeInvalid   BucketType = 0
	BktTypeCouchbase            = iota
	BktTypeMemcached            = iota
)

type VBucketState uint32

const (
	VBucketStateActive  = VBucketState(0x01)
	VBucketStateReplica = VBucketState(0x02)
	VBucketStatePending = VBucketState(0x03)
	VBucketStateDead    = VBucketState(0x04)
)
<commit_msg>Release v1.1.0<commit_after>package gocbcore

const (
	GoCbVersionStr = "v1.1.0"
)

type CommandMagic uint8

const (
	ReqMagic = CommandMagic(0x80)
	ResMagic = CommandMagic(0x81)
)

// CommandCode for memcached packets.
type CommandCode uint8

const (
	CmdGet                  = CommandCode(0x00)
	CmdSet                  = CommandCode(0x01)
	CmdAdd                  = CommandCode(0x02)
	CmdReplace              = CommandCode(0x03)
	CmdDelete               = CommandCode(0x04)
	CmdIncrement            = CommandCode(0x05)
	CmdDecrement            = CommandCode(0x06)
	CmdAppend               = CommandCode(0x0e)
	CmdPrepend              = CommandCode(0x0f)
	CmdStat                 = CommandCode(0x10)
	CmdTouch                = CommandCode(0x1c)
	CmdGAT                  = CommandCode(0x1d)
	CmdHello                = CommandCode(0x1f)
	CmdSASLListMechs        = CommandCode(0x20)
	CmdSASLAuth             = CommandCode(0x21)
	CmdSASLStep             = CommandCode(0x22)
	CmdGetAllVBSeqnos       = CommandCode(0x48)
	CmdDcpOpenConnection    = CommandCode(0x50)
	CmdDcpAddStream         = CommandCode(0x51)
	CmdDcpCloseStream       = CommandCode(0x52)
	CmdDcpStreamReq         = CommandCode(0x53)
	CmdDcpGetFailoverLog    = CommandCode(0x54)
	CmdDcpStreamEnd         = CommandCode(0x55)
	CmdDcpSnapshotMarker    = CommandCode(0x56)
	CmdDcpMutation          = CommandCode(0x57)
	CmdDcpDeletion          = CommandCode(0x58)
	CmdDcpExpiration        = CommandCode(0x59)
	CmdDcpFlush             = CommandCode(0x5a)
	CmdDcpSetVbucketState   = CommandCode(0x5b)
	CmdDcpNoop              = CommandCode(0x5c)
	CmdDcpBufferAck         = CommandCode(0x5d)
	CmdDcpControl           = CommandCode(0x5e)
	CmdGetReplica           = CommandCode(0x83)
	CmdSelectBucket         = CommandCode(0x89)
	CmdObserveSeqNo         = CommandCode(0x91)
	CmdObserve              = CommandCode(0x92)
	CmdGetLocked            = CommandCode(0x94)
	CmdUnlockKey            = CommandCode(0x95)
	CmdSetMeta              = CommandCode(0xa2)
	CmdDetMeta              = CommandCode(0xa8)
	CmdGetClusterConfig     = CommandCode(0xb5)
	CmdGetRandom            = CommandCode(0xb6)
	CmdSubDocGet            = CommandCode(0xc5)
	CmdSubDocExists         = CommandCode(0xc6)
	CmdSubDocDictAdd        = CommandCode(0xc7)
	CmdSubDocDictSet        = CommandCode(0xc8)
	CmdSubDocDelete         = CommandCode(0xc9)
	CmdSubDocReplace        = CommandCode(0xca)
	CmdSubDocArrayPushLast  = CommandCode(0xcb)
	CmdSubDocArrayPushFirst = CommandCode(0xcc)
	CmdSubDocArrayInsert    = CommandCode(0xcd)
	CmdSubDocArrayAddUnique = CommandCode(0xce)
	CmdSubDocCounter        = CommandCode(0xcf)
	CmdSubDocMultiLookup    = CommandCode(0xd0)
	CmdSubDocMultiMutation  = CommandCode(0xd1)
)

type SubDocFlag uint16

const (
	SubDocFlagMkDirP = SubDocFlag(0x01)
)

type SubDocOpType uint8

const (
	SubDocOpGet            = SubDocOpType(CmdSubDocGet)
	SubDocOpExists         = SubDocOpType(CmdSubDocExists)
	SubDocOpDictAdd        = SubDocOpType(CmdSubDocDictAdd)
	SubDocOpDictSet        = SubDocOpType(CmdSubDocDictSet)
	SubDocOpDelete         = SubDocOpType(CmdSubDocDelete)
	SubDocOpReplace        = SubDocOpType(CmdSubDocReplace)
	SubDocOpArrayPushLast  = SubDocOpType(CmdSubDocArrayPushLast)
	SubDocOpArrayPushFirst = SubDocOpType(CmdSubDocArrayPushFirst)
	SubDocOpArrayInsert    = SubDocOpType(CmdSubDocArrayInsert)
	SubDocOpArrayAddUnique = SubDocOpType(CmdSubDocArrayAddUnique)
	SubDocOpCounter        = SubDocOpType(CmdSubDocCounter)
)

type HelloFeature uint16

const (
	FeatureDatatype = HelloFeature(0x01)
	FeatureSeqNo    = HelloFeature(0x04)
)

// Status field for memcached response.
type StatusCode uint16

const (
	StatusSuccess            = StatusCode(0x00)
	StatusKeyNotFound        = StatusCode(0x01)
	StatusKeyExists          = StatusCode(0x02)
	StatusTooBig             = StatusCode(0x03)
	StatusInvalidArgs        = StatusCode(0x04)
	StatusNotStored          = StatusCode(0x05)
	StatusBadDelta           = StatusCode(0x06)
	StatusNotMyVBucket       = StatusCode(0x07)
	StatusNoBucket           = StatusCode(0x08)
	StatusAuthStale          = StatusCode(0x1f)
	StatusAuthError          = StatusCode(0x20)
	StatusAuthContinue       = StatusCode(0x21)
	StatusRangeError         = StatusCode(0x22)
	StatusRollback           = StatusCode(0x23)
	StatusAccessError        = StatusCode(0x24)
	StatusNotInitialized     = StatusCode(0x25)
	StatusUnknownCommand     = StatusCode(0x81)
	StatusOutOfMemory        = StatusCode(0x82)
	StatusNotSupported       = StatusCode(0x83)
	StatusInternalError      = StatusCode(0x84)
	StatusBusy               = StatusCode(0x85)
	StatusTmpFail            = StatusCode(0x86)
	StatusSubDocPathNotFound = StatusCode(0xc0)
	StatusSubDocPathMismatch = StatusCode(0xc1)
	StatusSubDocPathInvalid  = StatusCode(0xc2)
	StatusSubDocPathTooBig   = StatusCode(0xc3)
	StatusSubDocDocTooDeep   = StatusCode(0xc4)
	StatusSubDocCantInsert   = StatusCode(0xc5)
	StatusSubDocNotJson      = StatusCode(0xc6)
	StatusSubDocBadRange     = StatusCode(0xc7)
	StatusSubDocBadDelta     = StatusCode(0xc8)
	StatusSubDocPathExists   = StatusCode(0xc9)
	StatusSubDocValueTooDeep = StatusCode(0xca)
	StatusSubDocBadCombo     = StatusCode(0xcb)
	StatusSubDocBadMulti     = StatusCode(0xcc)
)

type KeyState uint8

const (
	KeyStateNotPersisted = KeyState(0x00)
	KeyStatePersisted    = KeyState(0x01)
	KeyStateNotFound     = KeyState(0x80)
	KeyStateDeleted      = KeyState(0x81)
)

type StreamEndStatus uint32

const (
	StreamEndOK           = StreamEndStatus(0x00)
	StreamEndClosed       = StreamEndStatus(0x01)
	StreamEndStateChanged = StreamEndStatus(0x02)
	StreamEndDisconnected = StreamEndStatus(0x03)
	StreamEndTooSlow      = StreamEndStatus(0x04)
)

type BucketType int

const (
	BktTypeInvalid   BucketType = 0
	BktTypeCouchbase            = iota
	BktTypeMemcached            = iota
)

type VBucketState uint32

const (
	VBucketStateActive  = VBucketState(0x01)
	VBucketStateReplica = VBucketState(0x02)
	VBucketStatePending = VBucketState(0x03)
	VBucketStateDead    = VBucketState(0x04)
)
<|endoftext|>


<commit_before>package mysql

import (
	"database/sql"
	"errors"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	lg "github.com/hiromaily/golibs/log"
	"github.com/hiromaily/golibs/times"
	u "github.com/hiromaily/golibs/utils"
	"reflect"
	"time"
)

//TODO:トランザクションの機能もあるので、どこかに追加しておく
//TODO:異なるlibraryを使っているが、各funcのInterfaceを統一すればよいのでは？
//http://qiita.com/tenntenn/items/dddb13c15643454a7c3b
//http://go-database-sql.org/

const tFomt = "2006-01-02 15:04:05"

type MS struct {
	DB         *sql.DB
	Rows       *sql.Rows
	Err        error
	ServerInfo //embeded
}

type ServerInfo struct {
	host   string
	port   uint16
	dbname string
	user   string
	pass   string
}

var dbInfo MS

//-----------------------------------------------------------------------------
// Basic
//-----------------------------------------------------------------------------
func New(host, dbname, user, pass string, port uint16) {
	var err error
	if dbInfo.DB == nil {
		dbInfo.host = host
		dbInfo.port = port
		dbInfo.dbname = dbname
		dbInfo.user = user
		dbInfo.pass = pass

		dbInfo.DB, err = dbInfo.Connection()
	}
	//lg.Debugf("dbInfo.db %+v\n", *dbInfo.DB)
	if err != nil {
		panic(err.Error())
	}
	return
}

// singleton architecture
func GetDBInstance() *MS {
	var err error
	if dbInfo.DB == nil {
		//TODO: it may be better to call New()
		dbInfo.DB, err = dbInfo.Connection()
	}
	if err != nil {
		panic(err.Error())
	}
	return &dbInfo
}

func (ms *MS) getDsn() string {
	//If use nil on Date column, set *time.Time
	//Be careful when parsing is required on Date type
	// e.g. db, err := sql.Open("mysql", "root:@/?parseTime=true")
	param := "?charset=utf8&parseTime=True&loc=Local"
	//user:password@tcp(localhost:3306)/dbname?tls=skip-verify&autocommit=true
	return fmt.Sprintf("%s:%s@tcp(%s:%d)/%s%s",
		ms.user, ms.pass, ms.host, ms.port, ms.dbname, param)
}

// Connection
// Be careful, sql.Open() doesn't return err. Use db.Ping() to check DB condition.
func (ms *MS) Connection() (*sql.DB, error) {
	//return sql.Open("mysql", getDsn())
	db, _ := sql.Open("mysql", ms.getDsn())
	return db, db.Ping()
}

// SetMaxIdleConns sets the maximum number of connections in the idle connection pool.
func (ms *MS) SetMaxIdleConns(n int) {
	ms.DB.SetMaxIdleConns(n)
}

//SetMaxOpenConns sets the maximum number of open connections to the database.
func (ms *MS) SetMaxOpenConns(n int) {
	ms.DB.SetMaxOpenConns(n)
}

// Close
func (ms *MS) Close() {
	ms.DB.Close()
}

//-----------------------------------------------------------------------------
// Select
//-----------------------------------------------------------------------------
// SELECT Count: Get number of rows
func (ms *MS) SelectCount(countSql string, args ...interface{}) (int, error) {
	//field on table
	var count int

	//1. create sql and exec
	//err := self.db.QueryRow("SELECT count(user_id) FROM t_users WHERE delete_flg=?", "0").Scan(&count)
	err := ms.DB.QueryRow(countSql, args...).Scan(&count)
	if err != nil {
		return 0, err
	}

	return count, nil
}

//get Rows and return db instance
func (ms *MS) SelectIns(selectSQL string, args ...interface{}) *MS {
	defer times.Track(time.Now(), "SelectIns()")
	//SelectSQLAllFieldIns() took 471.577µs

	//If no args, set nil

	//1. create sql and exec
	//rows, err := self.db.Query("SELECT * FROM t_users WHERE delete_flg=?", "0")
	ms.Rows, ms.Err = ms.DB.Query(selectSQL, args...)
	if ms.Err != nil {
		lg.Errorf("SelectSQLAllFieldIns()->ms.DB.Query():error is %s, \n %s", ms.Err, selectSQL)
	}

	return ms
}

//set extracted data into parameter variable
func (ms *MS) ScanOne(x interface{}) bool {
	//defer times.Track(time.Now(), "ScanOne()")
	//ScanOne() took 5.23µs

	if ms.Err != nil {
		lg.Errorf("ScanOne(): ms.Err has error: %s", ms.Err)
		return false
	}

	//e.g)v = person Person
	v := reflect.ValueOf(x)
	if v.Kind() != reflect.Ptr || v.IsNil() {
		ms.Err = errors.New("parameter is not valid. it sould be pointer and not nil.")
		return false
	} else {
		if v.Elem().Kind() == reflect.Struct {

			//create container to set scaned record on database
			values, scanArgs := makeScanArgs(v.Elem().Type())

			//check len(value) and column
			validateStructAndColumns(ms, values)
			if ms.Err != nil {
				return false
			}

			// rows.Next()
			ret := ms.Rows.Next()
			if !ret {
				//ms.Err = errors.New("nodata")
				return false
			}

			// rows.Scan()
			ms.Err = ms.Rows.Scan(scanArgs...)
			if ms.Err != nil {
				return false
			}

			//ms.Err = ms.Rows.Scan(v)
			scanStruct(values, v.Elem())
		} else {
			ms.Err = errors.New("parameter should be pointer of struct slice or struct")
			return false
		}
	}
	return true
}

func (ms *MS) Scan(x interface{}) bool {
	//defer times.Track(time.Now(), "Scan()")
	//Scan() took 465.971µs

	if ms.Err != nil {
		lg.Errorf("Scan(): ms.Err has error: %s", ms.Err)
		return false
	}

	//e.g)v = persons []Person
	v := reflect.ValueOf(x)

	if v.Kind() != reflect.Ptr || v.IsNil() {
		ms.Err = errors.New("parameter is not valid. it sould be pointer and not nil.")
		return false
	} else {
		if v.Elem().Kind() == reflect.Slice || v.Elem().Kind() == reflect.Array {
			elemType := v.Elem().Type().Elem() //reflects_test.TeacherInfo
			newElem := reflect.New(elemType).Elem()

			//create container to set scaned record on database
			values, scanArgs := makeScanArgs(newElem.Type())

			//check len(value) and column
			validateStructAndColumns(ms, values)
			if ms.Err != nil {
				return false
			}

			// rows.Next()
			cnt := 0
			for ms.Rows.Next() {
				ms.Err = ms.Rows.Scan(scanArgs...)
				if ms.Err != nil {
					return false
				}

				scanStruct(values, newElem)
				v.Elem().Set(reflect.Append(v.Elem(), newElem))
				cnt++
			}
			if cnt == 0 {
				return false
			}
		} else {
			ms.Err = errors.New("parameter is not valid. it sould be pointer and not nil.")
			return false
		}
	}

	return true
}

func makeScanArgs(structType reflect.Type) ([]interface{}, []interface{}) {
	values := make([]interface{}, structType.NumField())
	scanArgs := make([]interface{}, len(values))
	for i := range values {
		scanArgs[i] = &values[i]
	}

	return values, scanArgs
}

func validateStructAndColumns(ms *MS, values []interface{}) error {
	columns, err := ms.Rows.Columns()
	if err != nil {
		//when Rows are closed, error occur.
		ms.Err = err
		return ms.Err
	}
	if len(columns) != len(values) {
		ms.Err = fmt.Errorf("number of struct field(%d) doesn't match to columns of sql(%d).", len(values), len(columns))
		return ms.Err
	}
	return nil
}

//Set data
func scanStruct(values []interface{}, v reflect.Value) {
	structType := v.Type()
	for i := 0; i < structType.NumField(); i++ {
		val := reflect.ValueOf(values[i])
		switch val.Kind() {
		case reflect.Invalid:
			//nil: for now, it skips.
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			v.Field(i).Set(reflect.ValueOf(u.Itoi(values[i])))
		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
			v.Field(i).Set(reflect.ValueOf(u.ItoUi(values[i])))
		case reflect.Bool:
			v.Field(i).Set(reflect.ValueOf(u.Itob(values[i])))
		case reflect.String:
			v.Field(i).Set(reflect.ValueOf(u.Itos(values[i])))
		case reflect.Slice:
			if u.CheckInterface(values[i]) == "[]uint8" {
				v.Field(i).Set(reflect.ValueOf(u.ItoBS(values[i])))
			}
		//case reflect.Chan, reflect.Func, reflect.Ptr, reflect.Map:
		case reflect.Struct:
			//time.Time
			if u.CheckInterface(values[i]) == "time.Time" {
				v.Field(i).Set(reflect.ValueOf(u.ItoT(values[i]).Format(tFomt)))
			}
		default: // reflect.Array, reflect.Struct, reflect.Interface
			v.Field(i).Set(reflect.ValueOf(values[i]))
		}
	}
	return
}

// SELECT : Get All field you set(Though you get only record, use it.)
func (ms *MS) Select(selectSQL string, args ...interface{}) ([]map[string]interface{}, []string, error) {
	defer times.Track(time.Now(), "SelectSQLAllField()")
	//540.417µs

	//1. create sql and exec
	//rows, err := self.db.Query("SELECT * FROM t_users WHERE delete_flg=?", "0")
	rows, err := ms.DB.Query(selectSQL, args...)
	if err != nil {
		return nil, nil, err
	}

	return ms.convertRowsToMaps(rows)
}

// Convert result of select into Map[] type. Return multiple array map and interface(plural lines)
func (ms *MS) convertRowsToMaps(rows *sql.Rows) ([]map[string]interface{}, []string, error) {
	defer times.Track(time.Now(), "convertRowsToMaps()")
	//convertRowsToMaps() took 85.191µs

	// Get column name
	columns, err := rows.Columns()
	if err != nil {
		return nil, nil, err
	}

	values := make([]interface{}, len(columns))

	// rows.Scan は引数に `[]interface{}`が必要.
	scanArgs := make([]interface{}, len(values))

	for i := range values {
		scanArgs[i] = &values[i]
	}

	retMaps := []map[string]interface{}{}
	//
	for rows.Next() { //true or false
		//Get data into scanArgs
		err = rows.Scan(scanArgs...)

		if err != nil {
			return nil, columns, err
		}

		rowdata := map[string]interface{}{}

		//var v string
		for i, value := range values {
			if u.CheckInterface(value) == "[]uint8" {
				value = u.ItoBS(value)
			} else if u.CheckInterface(value) == "time.Time" {
				value = u.ItoT(value).Format(tFomt)
			}

			// Here we can check if the value is nil (NULL value)
			//if value == nil {
			//	v = "NULL"
			//} else {
			//	v = string(value)
			//}

			//if b, ok := value.([]byte); ok{
			//	v = string(b)
			//} else {
			//	v = "NULL"
			//}

			//rowdata[columns[i]] = v
			rowdata[columns[i]] = value
		}
		retMaps = append(retMaps, rowdata)
	}
	return retMaps, columns, nil
}

//-----------------------------------------------------------------------------
// Insert
//-----------------------------------------------------------------------------
func (self *MS) Insert(sql string, args ...interface{}) (int64, error) {
	//1.creates a prepared statement (placeholder)
	//insertSQL := "INSERT t_users SET first_name=?, last_name=?"
	stmt, err := self.DB.Prepare(sql)
	if err != nil {
		return 0, err
	}

	//2.set parameter to prepared statement
	//res, err := stmt.Exec("mitsuo", "fujita")
	res, err := stmt.Exec(args...)
	if err != nil {
		return 0, err
	}

	defer stmt.Close() //statementもcloseする必要がある

	//3.Get id from response
	//id, err := res.LastInsertId()
	return res.LastInsertId()
}

//-----------------------------------------------------------------------------
// UPDATE / DELETE
//-----------------------------------------------------------------------------
func (ms *MS) Exec(sql string, args ...interface{}) (int64, error) {

	//1.creates a prepared statement (placeholder)
	//updateSQL := "UPDATE t_users SET first_name=? WHERE user_id=?"
	stmt, err := ms.DB.Prepare(sql)
	if err != nil {
		return 0, err
	}

	//2.set parameter to prepared statement
	//res, err := stmt.Exec("genjiro", 3)
	res, err := stmt.Exec(args...)
	if err != nil {
		return 0, err
	}

	defer stmt.Close() //statementもcloseする必要がある

	//3.Get number of changed rows
	//rows, err := res.RowsAffected()
	return res.RowsAffected()
}

// Execution simply
func (ms *MS) Exec2(sql string, args ...interface{}) error {
	//result, err := self.db.Exec("INSERT t_users SET first_name=?, last_name=?", "Mika", "Haruda")
	_, err := ms.DB.Exec(sql, args...)
	return err
}
<commit_msg>added util func on db/mysql<commit_after>package mysql

import (
	"database/sql"
	"errors"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	lg "github.com/hiromaily/golibs/log"
	"github.com/hiromaily/golibs/times"
	u "github.com/hiromaily/golibs/utils"
	"reflect"
	"time"
)

//TODO:トランザクションの機能もあるので、どこかに追加しておく
//TODO:異なるlibraryを使っているが、各funcのInterfaceを統一すればよいのでは？
//http://qiita.com/tenntenn/items/dddb13c15643454a7c3b
//http://go-database-sql.org/

const tFomt = "2006-01-02 15:04:05"

type MS struct {
	DB         *sql.DB
	Rows       *sql.Rows
	Err        error
	ServerInfo //embeded
}

type ServerInfo struct {
	host   string
	port   uint16
	dbname string
	user   string
	pass   string
}

var dbInfo MS

//-----------------------------------------------------------------------------
// Basic
//-----------------------------------------------------------------------------
func New(host, dbname, user, pass string, port uint16) {
	var err error
	if dbInfo.DB == nil {
		dbInfo.host = host
		dbInfo.port = port
		dbInfo.dbname = dbname
		dbInfo.user = user
		dbInfo.pass = pass

		dbInfo.DB, err = dbInfo.Connection()
	}
	//lg.Debugf("dbInfo.db %+v\n", *dbInfo.DB)
	if err != nil {
		panic(err.Error())
	}
	return
}

// singleton architecture
func GetDBInstance() *MS {
	var err error
	if dbInfo.DB == nil {
		//TODO: it may be better to call New()
		dbInfo.DB, err = dbInfo.Connection()
	}
	if err != nil {
		panic(err.Error())
	}
	return &dbInfo
}

func (ms *MS) getDsn() string {
	//If use nil on Date column, set *time.Time
	//Be careful when parsing is required on Date type
	// e.g. db, err := sql.Open("mysql", "root:@/?parseTime=true")
	param := "?charset=utf8&parseTime=True&loc=Local"
	//user:password@tcp(localhost:3306)/dbname?tls=skip-verify&autocommit=true
	return fmt.Sprintf("%s:%s@tcp(%s:%d)/%s%s",
		ms.user, ms.pass, ms.host, ms.port, ms.dbname, param)
}

// Connection
// Be careful, sql.Open() doesn't return err. Use db.Ping() to check DB condition.
func (ms *MS) Connection() (*sql.DB, error) {
	//return sql.Open("mysql", getDsn())
	db, _ := sql.Open("mysql", ms.getDsn())
	return db, db.Ping()
}

// SetMaxIdleConns sets the maximum number of connections in the idle connection pool.
func (ms *MS) SetMaxIdleConns(n int) {
	ms.DB.SetMaxIdleConns(n)
}

//SetMaxOpenConns sets the maximum number of open connections to the database.
func (ms *MS) SetMaxOpenConns(n int) {
	ms.DB.SetMaxOpenConns(n)
}

// Close
func (ms *MS) Close() {
	ms.DB.Close()
}

//-----------------------------------------------------------------------------
// Select
//-----------------------------------------------------------------------------
// SELECT Count: Get number of rows
func (ms *MS) SelectCount(countSql string, args ...interface{}) (int, error) {
	//field on table
	var count int

	//1. create sql and exec
	//err := self.db.QueryRow("SELECT count(user_id) FROM t_users WHERE delete_flg=?", "0").Scan(&count)
	err := ms.DB.QueryRow(countSql, args...).Scan(&count)
	if err != nil {
		return 0, err
	}

	return count, nil
}

//get Rows and return db instance
func (ms *MS) SelectIns(selectSQL string, args ...interface{}) *MS {
	defer times.Track(time.Now(), "SelectIns()")
	//SelectSQLAllFieldIns() took 471.577µs

	//If no args, set nil

	//1. create sql and exec
	//rows, err := self.db.Query("SELECT * FROM t_users WHERE delete_flg=?", "0")
	ms.Rows, ms.Err = ms.DB.Query(selectSQL, args...)
	if ms.Err != nil {
		lg.Errorf("SelectSQLAllFieldIns()->ms.DB.Query():error is %s, \n %s", ms.Err, selectSQL)
	}

	return ms
}

//set extracted data into parameter variable
func (ms *MS) ScanOne(x interface{}) bool {
	//defer times.Track(time.Now(), "ScanOne()")
	//ScanOne() took 5.23µs

	if ms.Err != nil {
		lg.Errorf("ScanOne(): ms.Err has error: %s", ms.Err)
		return false
	}

	//e.g)v = person Person
	v := reflect.ValueOf(x)
	if v.Kind() != reflect.Ptr || v.IsNil() {
		ms.Err = errors.New("parameter is not valid. it sould be pointer and not nil.")
		return false
	} else {
		if v.Elem().Kind() == reflect.Struct {

			//create container to set scaned record on database
			values, scanArgs := makeScanArgs(v.Elem().Type())

			//check len(value) and column
			validateStructAndColumns(ms, values)
			if ms.Err != nil {
				return false
			}

			// rows.Next()
			ret := ms.Rows.Next()
			if !ret {
				//ms.Err = errors.New("nodata")
				return false
			}

			// rows.Scan()
			ms.Err = ms.Rows.Scan(scanArgs...)
			if ms.Err != nil {
				return false
			}

			//ms.Err = ms.Rows.Scan(v)
			scanStruct(values, v.Elem())
		} else {
			ms.Err = errors.New("parameter should be pointer of struct slice or struct")
			return false
		}
	}
	return true
}

func (ms *MS) Scan(x interface{}) bool {
	//defer times.Track(time.Now(), "Scan()")
	//Scan() took 465.971µs

	if ms.Err != nil {
		lg.Errorf("Scan(): ms.Err has error: %s", ms.Err)
		return false
	}

	//e.g)v = persons []Person
	v := reflect.ValueOf(x)

	if v.Kind() != reflect.Ptr || v.IsNil() {
		ms.Err = errors.New("parameter is not valid. it sould be pointer and not nil.")
		return false
	} else {
		if v.Elem().Kind() == reflect.Slice || v.Elem().Kind() == reflect.Array {
			elemType := v.Elem().Type().Elem() //reflects_test.TeacherInfo
			newElem := reflect.New(elemType).Elem()

			//create container to set scaned record on database
			values, scanArgs := makeScanArgs(newElem.Type())

			//check len(value) and column
			validateStructAndColumns(ms, values)
			if ms.Err != nil {
				return false
			}

			// rows.Next()
			cnt := 0
			for ms.Rows.Next() {
				ms.Err = ms.Rows.Scan(scanArgs...)
				if ms.Err != nil {
					return false
				}

				scanStruct(values, newElem)
				v.Elem().Set(reflect.Append(v.Elem(), newElem))
				cnt++
			}
			if cnt == 0 {
				return false
			}
		} else {
			ms.Err = errors.New("parameter is not valid. it sould be pointer and not nil.")
			return false
		}
	}

	return true
}

func makeScanArgs(structType reflect.Type) ([]interface{}, []interface{}) {
	values := make([]interface{}, structType.NumField())
	scanArgs := make([]interface{}, len(values))
	for i := range values {
		scanArgs[i] = &values[i]
	}

	return values, scanArgs
}

func validateStructAndColumns(ms *MS, values []interface{}) error {
	columns, err := ms.Rows.Columns()
	if err != nil {
		//when Rows are closed, error occur.
		ms.Err = err
		return ms.Err
	}
	if len(columns) != len(values) {
		ms.Err = fmt.Errorf("number of struct field(%d) doesn't match to columns of sql(%d).", len(values), len(columns))
		return ms.Err
	}
	return nil
}

//Set data
func scanStruct(values []interface{}, v reflect.Value) {
	structType := v.Type()
	for i := 0; i < structType.NumField(); i++ {
		val := reflect.ValueOf(values[i])
		switch val.Kind() {
		case reflect.Invalid:
			//nil: for now, it skips.
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			v.Field(i).Set(reflect.ValueOf(u.Itoi(values[i])))
		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
			v.Field(i).Set(reflect.ValueOf(u.ItoUi(values[i])))
		case reflect.Bool:
			v.Field(i).Set(reflect.ValueOf(u.Itob(values[i])))
		case reflect.String:
			v.Field(i).Set(reflect.ValueOf(u.Itos(values[i])))
		case reflect.Slice:
			if u.CheckInterface(values[i]) == "[]uint8" {
				v.Field(i).Set(reflect.ValueOf(u.ItoBS(values[i])))
			}
		//case reflect.Chan, reflect.Func, reflect.Ptr, reflect.Map:
		case reflect.Struct:
			//time.Time
			if u.CheckInterface(values[i]) == "time.Time" {
				v.Field(i).Set(reflect.ValueOf(u.ItoT(values[i]).Format(tFomt)))
			}
		default: // reflect.Array, reflect.Struct, reflect.Interface
			v.Field(i).Set(reflect.ValueOf(values[i]))
		}
	}
	return
}

// SELECT : Get All field you set(Though you get only record, use it.)
func (ms *MS) Select(selectSQL string, args ...interface{}) ([]map[string]interface{}, []string, error) {
	defer times.Track(time.Now(), "SelectSQLAllField()")
	//540.417µs

	//1. create sql and exec
	//rows, err := self.db.Query("SELECT * FROM t_users WHERE delete_flg=?", "0")
	rows, err := ms.DB.Query(selectSQL, args...)
	if err != nil {
		return nil, nil, err
	}

	return ms.convertRowsToMaps(rows)
}

// Convert result of select into Map[] type. Return multiple array map and interface(plural lines)
func (ms *MS) convertRowsToMaps(rows *sql.Rows) ([]map[string]interface{}, []string, error) {
	defer times.Track(time.Now(), "convertRowsToMaps()")
	//convertRowsToMaps() took 85.191µs

	// Get column name
	columns, err := rows.Columns()
	if err != nil {
		return nil, nil, err
	}

	values := make([]interface{}, len(columns))

	// rows.Scan は引数に `[]interface{}`が必要.
	scanArgs := make([]interface{}, len(values))

	for i := range values {
		scanArgs[i] = &values[i]
	}

	retMaps := []map[string]interface{}{}
	//
	for rows.Next() { //true or false
		//Get data into scanArgs
		err = rows.Scan(scanArgs...)

		if err != nil {
			return nil, columns, err
		}

		rowdata := map[string]interface{}{}

		//var v string
		for i, value := range values {
			if u.CheckInterface(value) == "[]uint8" {
				value = u.ItoBS(value)
			} else if u.CheckInterface(value) == "time.Time" {
				value = u.ItoT(value).Format(tFomt)
			}

			// Here we can check if the value is nil (NULL value)
			//if value == nil {
			//	v = "NULL"
			//} else {
			//	v = string(value)
			//}

			//if b, ok := value.([]byte); ok{
			//	v = string(b)
			//} else {
			//	v = "NULL"
			//}

			//rowdata[columns[i]] = v
			rowdata[columns[i]] = value
		}
		retMaps = append(retMaps, rowdata)
	}
	return retMaps, columns, nil
}

//-----------------------------------------------------------------------------
// Insert
//-----------------------------------------------------------------------------
func (self *MS) Insert(sql string, args ...interface{}) (int64, error) {
	//1.creates a prepared statement (placeholder)
	//insertSQL := "INSERT t_users SET first_name=?, last_name=?"
	stmt, err := self.DB.Prepare(sql)
	if err != nil {
		return 0, err
	}

	//2.set parameter to prepared statement
	//res, err := stmt.Exec("mitsuo", "fujita")
	res, err := stmt.Exec(args...)
	if err != nil {
		return 0, err
	}

	defer stmt.Close() //statementもcloseする必要がある

	//3.Get id from response
	//id, err := res.LastInsertId()
	return res.LastInsertId()
}

//-----------------------------------------------------------------------------
// UPDATE / DELETE
//-----------------------------------------------------------------------------
func (ms *MS) Exec(sql string, args ...interface{}) (int64, error) {

	//1.creates a prepared statement (placeholder)
	//updateSQL := "UPDATE t_users SET first_name=? WHERE user_id=?"
	stmt, err := ms.DB.Prepare(sql)
	if err != nil {
		return 0, err
	}

	//2.set parameter to prepared statement
	//res, err := stmt.Exec("genjiro", 3)
	res, err := stmt.Exec(args...)
	if err != nil {
		return 0, err
	}

	defer stmt.Close() //statementもcloseする必要がある

	//3.Get number of changed rows
	//rows, err := res.RowsAffected()
	return res.RowsAffected()
}

// Execution simply
func (ms *MS) Exec2(sql string, args ...interface{}) error {
	//result, err := self.db.Exec("INSERT t_users SET first_name=?, last_name=?", "Mika", "Haruda")
	_, err := ms.DB.Exec(sql, args...)
	return err
}

//-----------------------------------------------------------------------------
// Util
//-----------------------------------------------------------------------------
func ColumnForSQL(s interface{}) string {
	v := reflect.ValueOf(s)
	if v.Elem().Kind() == reflect.Slice || v.Elem().Kind() == reflect.Array {
		elemType := v.Elem().Type().Elem()
		newElem := reflect.New(elemType).Elem()
		return scanColumn(newElem)
	} else if v.Elem().Kind() == reflect.Struct {
		return scanColumn(v.Elem())
	}
	return ""
}

func scanColumn(val reflect.Value) string {
	var fieldName string

	for i := 0; i < val.NumField(); i++ {

		typeField := val.Type().Field(i)
		tag := typeField.Tag

		valid := tag.Get("column")

		if valid != "" {
			fieldName += valid + ","
		}
	}
	if fieldName != "" {
		//remove last comma
		fieldName = string(fieldName[:(len(fieldName) - 1)])
	}

	return fieldName
}
<|endoftext|>


<commit_before>package db

import (
	"fmt"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	// sql drivers
	_ "github.com/mattn/go-sqlite3"
)

func Database(sqls ...string) (*DB, error) {
	db, err := Connect(":memory:")
	if err != nil {
		return nil, err
	}

	if _, err := db.Setup(0); err != nil {
		db.Disconnect()
		return nil, err
	}

	for _, s := range sqls {
		err := db.exec(s)
		if err != nil {
			db.Disconnect()
			return nil, err
		}
	}

	return db, nil
}

var _ = Describe("Database Schema", func() {
	Describe("Initializing the schema", func() {
		Context("With a new database", func() {
			var db *DB

			BeforeEach(func() {
				var err error
				db, err = Connect(":memory:")

				Ω(err).ShouldNot(HaveOccurred())
				Ω(db.Connected()).Should(BeTrue())
			})

			It("should not create tables until Setup() is called", func() {
				Ω(db.exec("SELECT * FROM schema_info")).
					Should(HaveOccurred())
			})

			It("should create tables during Setup()", func() {
				_, err := db.Setup(0)
				Ω(err).ShouldNot(HaveOccurred())
				Ω(db.exec("SELECT * FROM schema_info")).
					Should(Succeed())
			})

			It("should set the version number in schema_info", func() {
				_, err := db.Setup(0)
				Ω(err).ShouldNot(HaveOccurred())

				r, err := db.query(`SELECT version FROM schema_info`)
				Ω(err).ShouldNot(HaveOccurred())
				Ω(r).ShouldNot(BeNil())
				Ω(r.Next()).Should(BeTrue())

				var v int
				Ω(r.Scan(&v)).Should(Succeed())
				Ω(v).Should(Equal(6))
			})

			It("creates the correct tables", func() {
				_, err := db.Setup(0)
				Ω(err).ShouldNot(HaveOccurred())

				tableExists := func(table string) {
					sql := fmt.Sprintf("SELECT * FROM %s", table)
					Ω(db.exec(sql)).Should(Succeed())
				}

				tableExists("targets")
				tableExists("stores")
				tableExists("jobs")
				tableExists("archives")
				tableExists("tasks")
			})
		})
	})

	Describe("Schema Version Interrogation", func() {
		It("should return an error for a bad database connection", func() {
			db, _ := Connect("/path/to/no/such/file")
			_, err := db.SchemaVersion()
			Ω(err).Should(HaveOccurred())
		})
	})
})
<commit_msg>Fix test assertions, re: latest schema<commit_after>package db

import (
	"fmt"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	// sql drivers
	_ "github.com/mattn/go-sqlite3"
)

func Database(sqls ...string) (*DB, error) {
	db, err := Connect(":memory:")
	if err != nil {
		return nil, err
	}

	if _, err := db.Setup(0); err != nil {
		db.Disconnect()
		return nil, err
	}

	for _, s := range sqls {
		err := db.exec(s)
		if err != nil {
			db.Disconnect()
			return nil, err
		}
	}

	return db, nil
}

var _ = Describe("Database Schema", func() {
	Describe("Initializing the schema", func() {
		Context("With a new database", func() {
			var db *DB

			BeforeEach(func() {
				var err error
				db, err = Connect(":memory:")

				Ω(err).ShouldNot(HaveOccurred())
				Ω(db.Connected()).Should(BeTrue())
			})

			It("should not create tables until Setup() is called", func() {
				Ω(db.exec("SELECT * FROM schema_info")).
					Should(HaveOccurred())
			})

			It("should create tables during Setup()", func() {
				_, err := db.Setup(0)
				Ω(err).ShouldNot(HaveOccurred())
				Ω(db.exec("SELECT * FROM schema_info")).
					Should(Succeed())
			})

			It("should set the version number in schema_info", func() {
				_, err := db.Setup(0)
				Ω(err).ShouldNot(HaveOccurred())

				r, err := db.query(`SELECT version FROM schema_info`)
				Ω(err).ShouldNot(HaveOccurred())
				Ω(r).ShouldNot(BeNil())
				Ω(r.Next()).Should(BeTrue())

				var v int
				Ω(r.Scan(&v)).Should(Succeed())
				Ω(v).Should(Equal(7))
			})

			It("creates the correct tables", func() {
				_, err := db.Setup(0)
				Ω(err).ShouldNot(HaveOccurred())

				tableExists := func(table string) {
					sql := fmt.Sprintf("SELECT * FROM %s", table)
					Ω(db.exec(sql)).Should(Succeed())
				}

				tableExists("targets")
				tableExists("stores")
				tableExists("jobs")
				tableExists("archives")
				tableExists("tasks")
			})
		})
	})

	Describe("Schema Version Interrogation", func() {
		It("should return an error for a bad database connection", func() {
			db, _ := Connect("/path/to/no/such/file")
			_, err := db.SchemaVersion()
			Ω(err).Should(HaveOccurred())
		})
	})
})
<|endoftext|>


<commit_before>package eveConsumer

import (
	"evedata/esi"
	"fmt"
	"log"

	"golang.org/x/net/context"
)

func (c *EVEConsumer) checkPublicStructures() {
	log.Printf("EVEConsumer: collecting structures")
	err := c.collectStructuresFromESI()
	if err != nil {
		log.Printf("EVEConsumer: collecting structures: %v", err)
	}
}

func (c *EVEConsumer) collectStructuresFromESI() error {
	r := struct {
		Value int
		Wait  int
	}{0, 0}

	if err := c.ctx.Db.Get(&r, `
		SELECT value, TIME_TO_SEC(TIMEDIFF(nextCheck, UTC_TIMESTAMP())) AS wait
			FROM states
			WHERE state = 'structures'
			LIMIT 1;
		`); err != nil {
		return err
	}

	if r.Wait >= 0 {
		return nil
	}

	w, _, err := c.ctx.ESI.UniverseApi.GetUniverseStructures(nil)
	if err != nil {
		return err
	}

	// Update state so we dont have two polling at once.
	_, err = c.ctx.Db.Exec("UPDATE states SET value = 1, nextCheck =? WHERE state = 'structures' LIMIT 1")

	for _, s := range w {
		c.updateStructure(s)
	}
	return nil
}

func (c *EVEConsumer) updateStructure(s int64) error {
	ctx := context.WithValue(context.TODO(), esi.ContextAuth, c.ctx.ESIPublicToken)
	struc, _, err := c.ctx.ESI.UniverseApi.GetUniverseStructuresStructureId(ctx, s, nil)
	fmt.Printf("%+v\n", struc)
	if err != nil {
		return err
	}

	_, err = c.ctx.Db.Exec(`INSERT INTO staStations
					(stationID, solarSystemID, stationName, x, y, z, constellationID, regionID)
					VALUES(?,?,?,?,?,?,constellationIDBySolarSystem(solarSystemID),regionIDBySolarSystem(solarSystemID))
					ON DUPLICATE KEY UPDATE stationName=VALUES(stationName),solarSystemID=VALUES(solarSystemID),
					x=VALUES(x),y=VALUES(y),z=VALUES(z),constellationID=constellationIDBySolarSystem(VALUES(solarSystemID)),regionID=regionIDBySolarSystem(VALUES(solarSystemID));`,
		s, struc.SolarSystemId, struc.Name, struc.Position.X, struc.Position.Y, struc.Position.Z)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>remove print message<commit_after>package eveConsumer

import (
	"evedata/esi"
	"log"

	"golang.org/x/net/context"
)

func (c *EVEConsumer) checkPublicStructures() {
	log.Printf("EVEConsumer: collecting structures")
	err := c.collectStructuresFromESI()
	if err != nil {
		log.Printf("EVEConsumer: collecting structures: %v", err)
	}
}

func (c *EVEConsumer) collectStructuresFromESI() error {
	r := struct {
		Value int
		Wait  int
	}{0, 0}

	if err := c.ctx.Db.Get(&r, `
		SELECT value, TIME_TO_SEC(TIMEDIFF(nextCheck, UTC_TIMESTAMP())) AS wait
			FROM states
			WHERE state = 'structures'
			LIMIT 1;
		`); err != nil {
		return err
	}

	if r.Wait >= 0 {
		return nil
	}

	w, _, err := c.ctx.ESI.UniverseApi.GetUniverseStructures(nil)
	if err != nil {
		return err
	}

	// Update state so we dont have two polling at once.
	_, err = c.ctx.Db.Exec("UPDATE states SET value = 1, nextCheck =? WHERE state = 'structures' LIMIT 1")

	for _, s := range w {
		c.updateStructure(s)
	}
	return nil
}

func (c *EVEConsumer) updateStructure(s int64) error {
	ctx := context.WithValue(context.TODO(), esi.ContextAuth, c.ctx.ESIPublicToken)
	struc, _, err := c.ctx.ESI.UniverseApi.GetUniverseStructuresStructureId(ctx, s, nil)
	if err != nil {
		return err
	}

	_, err = c.ctx.Db.Exec(`INSERT INTO staStations
					(stationID, solarSystemID, stationName, x, y, z, constellationID, regionID)
					VALUES(?,?,?,?,?,?,constellationIDBySolarSystem(solarSystemID),regionIDBySolarSystem(solarSystemID))
					ON DUPLICATE KEY UPDATE stationName=VALUES(stationName),solarSystemID=VALUES(solarSystemID),
					x=VALUES(x),y=VALUES(y),z=VALUES(z),constellationID=constellationIDBySolarSystem(VALUES(solarSystemID)),regionID=regionIDBySolarSystem(VALUES(solarSystemID));`,
		s, struc.SolarSystemId, struc.Name, struc.Position.X, struc.Position.Y, struc.Position.Z)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package gqltesting

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"sort"
	"strconv"
	"testing"

	graphql "github.com/graph-gophers/graphql-go"
	"github.com/graph-gophers/graphql-go/errors"
)

// Test is a GraphQL test case to be used with RunTest(s).
type Test struct {
	Context        context.Context
	Schema         *graphql.Schema
	Query          string
	OperationName  string
	Variables      map[string]interface{}
	ExpectedResult string
	ExpectedErrors []*errors.QueryError
	RawResponse    bool
}

// RunTests runs the given GraphQL test cases as subtests.
func RunTests(t *testing.T, tests []*Test) {
	if len(tests) == 1 {
		RunTest(t, tests[0])
		return
	}

	for i, test := range tests {
		t.Run(strconv.Itoa(i+1), func(t *testing.T) {
			RunTest(t, test)
		})
	}
}

// RunTest runs a single GraphQL test case.
func RunTest(t *testing.T, test *Test) {
	if test.Context == nil {
		test.Context = context.Background()
	}
	result := test.Schema.Exec(test.Context, test.Query, test.OperationName, test.Variables)

	checkErrors(t, test.ExpectedErrors, result.Errors)

	if test.ExpectedResult == "" {
		if result.Data != nil {
			t.Fatalf("got: %s", result.Data)
			t.Fatalf("want: null")
		}
		return
	}

	// Verify JSON to avoid red herring errors.
	var got []byte

	if test.RawResponse {
		value, err := result.Data.MarshalJSON()
		if err != nil {
			t.Fatalf("got: unable to marshal JSON response: %s", err)
		}
		got = value
	} else {
		value, err := formatJSON(result.Data)
		if err != nil {
			t.Fatalf("got: invalid JSON: %s", err)
		}
		got = value
	}

	want, err := formatJSON([]byte(test.ExpectedResult))
	if err != nil {
		t.Fatalf("want: invalid JSON: %s", err)
	}

	if !bytes.Equal(got, want) {
		t.Logf("got:  %s", got)
		t.Logf("want: %s", want)
		t.Fail()
	}
}

func formatJSON(data []byte) ([]byte, error) {
	var v interface{}
	if err := json.Unmarshal(data, &v); err != nil {
		return nil, err
	}
	formatted, err := json.Marshal(v)
	if err != nil {
		return nil, err
	}
	return formatted, nil
}

func checkErrors(t *testing.T, want, got []*errors.QueryError) {
	sortErrors(want)
	sortErrors(got)

	// Clear the underlying error before the DeepEqual check.  It's too
	// much to ask the tester to include the raw failing error.
	for _, err := range got {
		err.Err = nil
	}

	if !reflect.DeepEqual(got, want) {
		t.Fatalf("unexpected error: got %+v, want %+v", got, want)
	}
}

func sortErrors(errors []*errors.QueryError) {
	if len(errors) <= 1 {
		return
	}
	sort.Slice(errors, func(i, j int) bool {
		return fmt.Sprintf("%s", errors[i].Path) < fmt.Sprintf("%s", errors[j].Path)
	})
}
<commit_msg>fix minor unreachable code caused by t.Fatalf<commit_after>package gqltesting

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"sort"
	"strconv"
	"testing"

	graphql "github.com/graph-gophers/graphql-go"
	"github.com/graph-gophers/graphql-go/errors"
)

// Test is a GraphQL test case to be used with RunTest(s).
type Test struct {
	Context        context.Context
	Schema         *graphql.Schema
	Query          string
	OperationName  string
	Variables      map[string]interface{}
	ExpectedResult string
	ExpectedErrors []*errors.QueryError
	RawResponse    bool
}

// RunTests runs the given GraphQL test cases as subtests.
func RunTests(t *testing.T, tests []*Test) {
	if len(tests) == 1 {
		RunTest(t, tests[0])
		return
	}

	for i, test := range tests {
		t.Run(strconv.Itoa(i+1), func(t *testing.T) {
			RunTest(t, test)
		})
	}
}

// RunTest runs a single GraphQL test case.
func RunTest(t *testing.T, test *Test) {
	if test.Context == nil {
		test.Context = context.Background()
	}
	result := test.Schema.Exec(test.Context, test.Query, test.OperationName, test.Variables)

	checkErrors(t, test.ExpectedErrors, result.Errors)

	if test.ExpectedResult == "" {
		if result.Data != nil {
			t.Fatalf("got: %s, want: null", result.Data)
		}
		return
	}

	// Verify JSON to avoid red herring errors.
	var got []byte

	if test.RawResponse {
		value, err := result.Data.MarshalJSON()
		if err != nil {
			t.Fatalf("got: unable to marshal JSON response: %s", err)
		}
		got = value
	} else {
		value, err := formatJSON(result.Data)
		if err != nil {
			t.Fatalf("got: invalid JSON: %s", err)
		}
		got = value
	}

	want, err := formatJSON([]byte(test.ExpectedResult))
	if err != nil {
		t.Fatalf("want: invalid JSON: %s", err)
	}

	if !bytes.Equal(got, want) {
		t.Logf("got:  %s", got)
		t.Logf("want: %s", want)
		t.Fail()
	}
}

func formatJSON(data []byte) ([]byte, error) {
	var v interface{}
	if err := json.Unmarshal(data, &v); err != nil {
		return nil, err
	}
	formatted, err := json.Marshal(v)
	if err != nil {
		return nil, err
	}
	return formatted, nil
}

func checkErrors(t *testing.T, want, got []*errors.QueryError) {
	sortErrors(want)
	sortErrors(got)

	// Clear the underlying error before the DeepEqual check.  It's too
	// much to ask the tester to include the raw failing error.
	for _, err := range got {
		err.Err = nil
	}

	if !reflect.DeepEqual(got, want) {
		t.Fatalf("unexpected error: got %+v, want %+v", got, want)
	}
}

func sortErrors(errors []*errors.QueryError) {
	if len(errors) <= 1 {
		return
	}
	sort.Slice(errors, func(i, j int) bool {
		return fmt.Sprintf("%s", errors[i].Path) < fmt.Sprintf("%s", errors[j].Path)
	})
}
<|endoftext|>


<commit_before>package file

import (
	"io"
	"syscall"
)

type lockCloser struct {
	fd syscall.Handlle
}

func (l lockCloser) Close() error {
	return syscall.Close(l.fd)
}

func (osFileSystem) Lock(name string) (io.Closer, error) {
	path, err := syscall.UTF16PtrFromString(name)
	if err != nil {
		return nil, err
	}
	fd, err := syscall.CreateFile(
		path,
		syscall.GENERIC_READ|syscall.GENERIC_WRITE,
		0,
		nil,
		syscall.OPEN_ALWAYS,
		syscall.FILE_ATTRIBUTE_NORMAL,
		0,
	)
	if err != nil {
		return nil, err
	}
	return lockCloser{fd: fd}
}
<commit_msg>Fix file lock on Windows<commit_after>package file

import (
	"io"
	"syscall"
)

type lockCloser struct {
	fd syscall.Handle
}

func (l lockCloser) Close() error {
	return syscall.Close(l.fd)
}

func (osFileSystem) Lock(name string) (io.Closer, error) {
	path, err := syscall.UTF16PtrFromString(name)
	if err != nil {
		return nil, err
	}
	fd, err := syscall.CreateFile(
		path,
		syscall.GENERIC_READ|syscall.GENERIC_WRITE,
		0,
		nil,
		syscall.OPEN_ALWAYS,
		syscall.FILE_ATTRIBUTE_NORMAL,
		0,
	)
	if err != nil {
		return nil, err
	}
	return lockCloser{fd: fd}, nil
}
<|endoftext|>


<commit_before>package specrunner

import (
	"fmt"
	"os"
	"os/signal"
	"sync"
	"github.com/onsi/ginkgo/config"
	"github.com/onsi/ginkgo/internal/leafnodes"
	"github.com/onsi/ginkgo/internal/spec"
	Writer "github.com/onsi/ginkgo/internal/writer"
	"github.com/onsi/ginkgo/reporters"
	"github.com/onsi/ginkgo/types"

	"time"
)

type SpecRunner struct {
	description     string
	beforeSuiteNode leafnodes.SuiteNode
	specs           *spec.Specs
	afterSuiteNode  leafnodes.SuiteNode
	reporters       []reporters.Reporter
	startTime       time.Time
	suiteID         string
	runningSpec     *spec.Spec
	writer          Writer.WriterInterface
	config          config.GinkgoConfigType
	interrupted     bool
	lock            *sync.Mutex
}

func New(description string, beforeSuiteNode leafnodes.SuiteNode, specs *spec.Specs, afterSuiteNode leafnodes.SuiteNode, reporters []reporters.Reporter, writer Writer.WriterInterface, config config.GinkgoConfigType) *SpecRunner {
	return &SpecRunner{
		description:     description,
		beforeSuiteNode: beforeSuiteNode,
		specs:           specs,
		afterSuiteNode:  afterSuiteNode,
		reporters:       reporters,
		writer:          writer,
		config:          config,
		suiteID:         randomID(),
		lock:            &sync.Mutex{},
	}
}

func (runner *SpecRunner) Run() bool {
	runner.reportSuiteWillBegin()
	go runner.registerForInterrupts()

	suitePassed := runner.runBeforeSuite()

	if suitePassed {
		suitePassed = runner.runSpecs()
	}

	runner.blockForeverIfInterrupted()

	suitePassed = runner.runAfterSuite() && suitePassed

	runner.reportSuiteDidEnd()

	return suitePassed
}

func (runner *SpecRunner) runBeforeSuite() bool {
	if runner.beforeSuiteNode == nil || runner.wasInterrupted() {
		return true
	}

	runner.writer.Truncate()
	conf := runner.config
	passed := runner.beforeSuiteNode.Run(conf.ParallelNode, conf.ParallelTotal, conf.SyncHost)
	if !passed {
		runner.writer.DumpOut()
	}
	runner.reportBeforeSuite(runner.beforeSuiteNode.Summary())
	return passed
}

func (runner *SpecRunner) runAfterSuite() bool {
	if runner.afterSuiteNode == nil {
		return true
	}

	runner.writer.Truncate()
	conf := runner.config
	passed := runner.afterSuiteNode.Run(conf.ParallelNode, conf.ParallelTotal, conf.SyncHost)
	if !passed {
		runner.writer.DumpOut()
	}
	runner.reportAfterSuite(runner.afterSuiteNode.Summary())
	return passed
}

func (runner *SpecRunner) runSpecs() bool {
	suiteFailed := false
	for _, spec := range runner.specs.Specs() {
		if runner.wasInterrupted() {
			return suiteFailed
		}
		runner.writer.Truncate()

		runner.reportSpecWillRun(spec)

		if !spec.Skipped() && !spec.Pending() {
			runner.runningSpec = spec
			spec.Run()
			runner.runningSpec = nil
			if spec.Failed() {
				suiteFailed = true
				runner.writer.DumpOut()
			}
		} else if spec.Pending() && runner.config.FailOnPending {
			suiteFailed = true
		}

		runner.reportSpecDidComplete(spec)
	}

	return !suiteFailed
}

func (runner *SpecRunner) CurrentSpecSummary() (*types.SpecSummary, bool) {
	if runner.runningSpec == nil {
		return nil, false
	}

	return runner.runningSpec.Summary(runner.suiteID), true
}

func (runner *SpecRunner) registerForInterrupts() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	<-c
	signal.Stop(c)
	runner.markInterrupted()
	go runner.registerForHardInterrupts()
	if runner.afterSuiteNode != nil {
		fmt.Fprintln(os.Stderr, "\nReceived interrupt.  Running AfterSuite...\n^C again to terminate immediately")
		runner.runAfterSuite()
	}
	runner.reportSuiteDidEnd()
	os.Exit(1)
}

func (runner *SpecRunner) registerForHardInterrupts() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	<-c
	fmt.Fprintln(os.Stderr, "\nReceived second interrupt.  Shutting down.")
	os.Exit(1)
}

func (runner *SpecRunner) blockForeverIfInterrupted() {
	runner.lock.Lock()
	interrupted := runner.interrupted
	runner.lock.Unlock()

	if interrupted {
		select {}
	}
}

func (runner *SpecRunner) markInterrupted() {
	runner.lock.Lock()
	defer runner.lock.Unlock()
	runner.interrupted = true
}

func (runner *SpecRunner) wasInterrupted() bool {
	runner.lock.Lock()
	defer runner.lock.Unlock()
	return runner.interrupted
}

func (runner *SpecRunner) reportSuiteWillBegin() {
	runner.startTime = time.Now()
	summary := runner.summary()
	for _, reporter := range runner.reporters {
		reporter.SpecSuiteWillBegin(runner.config, summary)
	}
}

func (runner *SpecRunner) reportBeforeSuite(summary *types.SetupSummary) {
	for _, reporter := range runner.reporters {
		reporter.BeforeSuiteDidRun(summary)
	}
}

func (runner *SpecRunner) reportAfterSuite(summary *types.SetupSummary) {
	for _, reporter := range runner.reporters {
		reporter.AfterSuiteDidRun(summary)
	}
}

func (runner *SpecRunner) reportSpecWillRun(spec *spec.Spec) {
	summary := spec.Summary(runner.suiteID)
	for _, reporter := range runner.reporters {
		reporter.SpecWillRun(summary)
	}
}

func (runner *SpecRunner) reportSpecDidComplete(spec *spec.Spec) {
	summary := spec.Summary(runner.suiteID)
	for _, reporter := range runner.reporters {
		reporter.SpecDidComplete(summary)
	}
}

func (runner *SpecRunner) reportSuiteDidEnd() {
	summary := runner.summary()
	summary.RunTime = time.Since(runner.startTime)
	for _, reporter := range runner.reporters {
		reporter.SpecSuiteDidEnd(summary)
	}
}

func (runner *SpecRunner) countSpecsSatisfying(filter func(ex *spec.Spec) bool) (count int) {
	count = 0

	for _, spec := range runner.specs.Specs() {
		if filter(spec) {
			count++
		}
	}

	return count
}

func (runner *SpecRunner) summary() *types.SuiteSummary {
	numberOfSpecsThatWillBeRun := runner.countSpecsSatisfying(func(ex *spec.Spec) bool {
		return !ex.Skipped() && !ex.Pending()
	})

	numberOfPendingSpecs := runner.countSpecsSatisfying(func(ex *spec.Spec) bool {
		return ex.Pending()
	})

	numberOfSkippedSpecs := runner.countSpecsSatisfying(func(ex *spec.Spec) bool {
		return ex.Skipped()
	})

	numberOfPassedSpecs := runner.countSpecsSatisfying(func(ex *spec.Spec) bool {
		return ex.Passed()
	})

	numberOfFailedSpecs := runner.countSpecsSatisfying(func(ex *spec.Spec) bool {
		return ex.Failed()
	})

	success := true

	if numberOfFailedSpecs > 0 {
		success = false
	} else if numberOfPendingSpecs > 0 && runner.config.FailOnPending {
		success = false
	} else if runner.beforeSuiteNode != nil && !runner.beforeSuiteNode.Passed() {
		success = false
		numberOfFailedSpecs = numberOfSpecsThatWillBeRun
	} else if runner.afterSuiteNode != nil && !runner.afterSuiteNode.Passed() {
		success = false
	} else if runner.wasInterrupted() {
		success = false
	}

	return &types.SuiteSummary{
		SuiteDescription: runner.description,
		SuiteSucceeded:   success,
		SuiteID:          runner.suiteID,

		NumberOfSpecsBeforeParallelization: runner.specs.NumberOfOriginalSpecs(),
		NumberOfTotalSpecs:                 len(runner.specs.Specs()),
		NumberOfSpecsThatWillBeRun:         numberOfSpecsThatWillBeRun,
		NumberOfPendingSpecs:               numberOfPendingSpecs,
		NumberOfSkippedSpecs:               numberOfSkippedSpecs,
		NumberOfPassedSpecs:                numberOfPassedSpecs,
		NumberOfFailedSpecs:                numberOfFailedSpecs,
	}
}
<commit_msg>handle SIGTERM as well as SIGINT<commit_after>package specrunner

import (
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"github.com/onsi/ginkgo/config"
	"github.com/onsi/ginkgo/internal/leafnodes"
	"github.com/onsi/ginkgo/internal/spec"
	Writer "github.com/onsi/ginkgo/internal/writer"
	"github.com/onsi/ginkgo/reporters"
	"github.com/onsi/ginkgo/types"

	"time"
)

type SpecRunner struct {
	description     string
	beforeSuiteNode leafnodes.SuiteNode
	specs           *spec.Specs
	afterSuiteNode  leafnodes.SuiteNode
	reporters       []reporters.Reporter
	startTime       time.Time
	suiteID         string
	runningSpec     *spec.Spec
	writer          Writer.WriterInterface
	config          config.GinkgoConfigType
	interrupted     bool
	lock            *sync.Mutex
}

func New(description string, beforeSuiteNode leafnodes.SuiteNode, specs *spec.Specs, afterSuiteNode leafnodes.SuiteNode, reporters []reporters.Reporter, writer Writer.WriterInterface, config config.GinkgoConfigType) *SpecRunner {
	return &SpecRunner{
		description:     description,
		beforeSuiteNode: beforeSuiteNode,
		specs:           specs,
		afterSuiteNode:  afterSuiteNode,
		reporters:       reporters,
		writer:          writer,
		config:          config,
		suiteID:         randomID(),
		lock:            &sync.Mutex{},
	}
}

func (runner *SpecRunner) Run() bool {
	runner.reportSuiteWillBegin()
	go runner.registerForInterrupts()

	suitePassed := runner.runBeforeSuite()

	if suitePassed {
		suitePassed = runner.runSpecs()
	}

	runner.blockForeverIfInterrupted()

	suitePassed = runner.runAfterSuite() && suitePassed

	runner.reportSuiteDidEnd()

	return suitePassed
}

func (runner *SpecRunner) runBeforeSuite() bool {
	if runner.beforeSuiteNode == nil || runner.wasInterrupted() {
		return true
	}

	runner.writer.Truncate()
	conf := runner.config
	passed := runner.beforeSuiteNode.Run(conf.ParallelNode, conf.ParallelTotal, conf.SyncHost)
	if !passed {
		runner.writer.DumpOut()
	}
	runner.reportBeforeSuite(runner.beforeSuiteNode.Summary())
	return passed
}

func (runner *SpecRunner) runAfterSuite() bool {
	if runner.afterSuiteNode == nil {
		return true
	}

	runner.writer.Truncate()
	conf := runner.config
	passed := runner.afterSuiteNode.Run(conf.ParallelNode, conf.ParallelTotal, conf.SyncHost)
	if !passed {
		runner.writer.DumpOut()
	}
	runner.reportAfterSuite(runner.afterSuiteNode.Summary())
	return passed
}

func (runner *SpecRunner) runSpecs() bool {
	suiteFailed := false
	for _, spec := range runner.specs.Specs() {
		if runner.wasInterrupted() {
			return suiteFailed
		}
		runner.writer.Truncate()

		runner.reportSpecWillRun(spec)

		if !spec.Skipped() && !spec.Pending() {
			runner.runningSpec = spec
			spec.Run()
			runner.runningSpec = nil
			if spec.Failed() {
				suiteFailed = true
				runner.writer.DumpOut()
			}
		} else if spec.Pending() && runner.config.FailOnPending {
			suiteFailed = true
		}

		runner.reportSpecDidComplete(spec)
	}

	return !suiteFailed
}

func (runner *SpecRunner) CurrentSpecSummary() (*types.SpecSummary, bool) {
	if runner.runningSpec == nil {
		return nil, false
	}

	return runner.runningSpec.Summary(runner.suiteID), true
}

func (runner *SpecRunner) registerForInterrupts() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	<-c
	signal.Stop(c)
	runner.markInterrupted()
	go runner.registerForHardInterrupts()
	if runner.afterSuiteNode != nil {
		fmt.Fprintln(os.Stderr, "\nReceived interrupt.  Running AfterSuite...\n^C again to terminate immediately")
		runner.runAfterSuite()
	}
	runner.reportSuiteDidEnd()
	os.Exit(1)
}

func (runner *SpecRunner) registerForHardInterrupts() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)

	<-c
	fmt.Fprintln(os.Stderr, "\nReceived second interrupt.  Shutting down.")
	os.Exit(1)
}

func (runner *SpecRunner) blockForeverIfInterrupted() {
	runner.lock.Lock()
	interrupted := runner.interrupted
	runner.lock.Unlock()

	if interrupted {
		select {}
	}
}

func (runner *SpecRunner) markInterrupted() {
	runner.lock.Lock()
	defer runner.lock.Unlock()
	runner.interrupted = true
}

func (runner *SpecRunner) wasInterrupted() bool {
	runner.lock.Lock()
	defer runner.lock.Unlock()
	return runner.interrupted
}

func (runner *SpecRunner) reportSuiteWillBegin() {
	runner.startTime = time.Now()
	summary := runner.summary()
	for _, reporter := range runner.reporters {
		reporter.SpecSuiteWillBegin(runner.config, summary)
	}
}

func (runner *SpecRunner) reportBeforeSuite(summary *types.SetupSummary) {
	for _, reporter := range runner.reporters {
		reporter.BeforeSuiteDidRun(summary)
	}
}

func (runner *SpecRunner) reportAfterSuite(summary *types.SetupSummary) {
	for _, reporter := range runner.reporters {
		reporter.AfterSuiteDidRun(summary)
	}
}

func (runner *SpecRunner) reportSpecWillRun(spec *spec.Spec) {
	summary := spec.Summary(runner.suiteID)
	for _, reporter := range runner.reporters {
		reporter.SpecWillRun(summary)
	}
}

func (runner *SpecRunner) reportSpecDidComplete(spec *spec.Spec) {
	summary := spec.Summary(runner.suiteID)
	for _, reporter := range runner.reporters {
		reporter.SpecDidComplete(summary)
	}
}

func (runner *SpecRunner) reportSuiteDidEnd() {
	summary := runner.summary()
	summary.RunTime = time.Since(runner.startTime)
	for _, reporter := range runner.reporters {
		reporter.SpecSuiteDidEnd(summary)
	}
}

func (runner *SpecRunner) countSpecsSatisfying(filter func(ex *spec.Spec) bool) (count int) {
	count = 0

	for _, spec := range runner.specs.Specs() {
		if filter(spec) {
			count++
		}
	}

	return count
}

func (runner *SpecRunner) summary() *types.SuiteSummary {
	numberOfSpecsThatWillBeRun := runner.countSpecsSatisfying(func(ex *spec.Spec) bool {
		return !ex.Skipped() && !ex.Pending()
	})

	numberOfPendingSpecs := runner.countSpecsSatisfying(func(ex *spec.Spec) bool {
		return ex.Pending()
	})

	numberOfSkippedSpecs := runner.countSpecsSatisfying(func(ex *spec.Spec) bool {
		return ex.Skipped()
	})

	numberOfPassedSpecs := runner.countSpecsSatisfying(func(ex *spec.Spec) bool {
		return ex.Passed()
	})

	numberOfFailedSpecs := runner.countSpecsSatisfying(func(ex *spec.Spec) bool {
		return ex.Failed()
	})

	success := true

	if numberOfFailedSpecs > 0 {
		success = false
	} else if numberOfPendingSpecs > 0 && runner.config.FailOnPending {
		success = false
	} else if runner.beforeSuiteNode != nil && !runner.beforeSuiteNode.Passed() {
		success = false
		numberOfFailedSpecs = numberOfSpecsThatWillBeRun
	} else if runner.afterSuiteNode != nil && !runner.afterSuiteNode.Passed() {
		success = false
	} else if runner.wasInterrupted() {
		success = false
	}

	return &types.SuiteSummary{
		SuiteDescription: runner.description,
		SuiteSucceeded:   success,
		SuiteID:          runner.suiteID,

		NumberOfSpecsBeforeParallelization: runner.specs.NumberOfOriginalSpecs(),
		NumberOfTotalSpecs:                 len(runner.specs.Specs()),
		NumberOfSpecsThatWillBeRun:         numberOfSpecsThatWillBeRun,
		NumberOfPendingSpecs:               numberOfPendingSpecs,
		NumberOfSkippedSpecs:               numberOfSkippedSpecs,
		NumberOfPassedSpecs:                numberOfPassedSpecs,
		NumberOfFailedSpecs:                numberOfFailedSpecs,
	}
}
<|endoftext|>