<commit_before>// Copyright 2010 The GoGo Authors. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

// This file holds the basic scanning routines that separate a source file
// into the various tokens

package main

import "./libgogo/_obj/libgogo"

// Token struct holding the relevant data of a parsed token.
type Token struct {
    id uint64; // The id. Is one of TOKEN_*
    intValue uint64; // value storing the integer value if the token is TOKEN_INTEGER
    strValue string; // Value storing the token string if the token is TOKEN_STRING or TOKEN_IDENTIFIER
    nextChar byte; // Sometime the next char is already read. It is stored here to be re-assigned in the next GetNextToken() round
};

func GetNextTokenRaw(fd uint64, tok *Token) {
    var singleChar byte; // Byte holding the last read value
    // Flag indicating whether we are in a comment.
    // 0 for no comment
    // 1 for a single line comment
    // 2 for a multi line comment
    var inComment uint64;
    var done uint64; // Flag indicating whether a cycle (Token) is finsihed
    var spaceDone uint64; // Flag indicating whether an abolishment cycle is finished

    // Initialize variables
    done = 0;
    spaceDone = 0;
    inComment = 0;

    tok.strValue = "";

    // If the previous cycle had to read the next char (and stored it), it is
    // now used as first read
    if tok.nextChar == 0 {
        singleChar = libgogo.GetChar(fd)
    } else {
        singleChar = tok.nextChar;
        tok.nextChar = 0;
    }

    // check if it is a valid read, or an EOF
    if singleChar == 0 {
        tok.id = TOKEN_EOS;
        done = 1;
        spaceDone = 1;
    }

    //
    // Cleaning Tasks
    // The next part strips out spaces, newlines, tabs, and comments
    // Comments can either be single line with double slashes (//) or multiline
    // using C++ syntax /* */
    //
    for ; spaceDone != 1; {

        // check whether a comment is starting
        if singleChar == '/' {
            // if we are in a comment skip the rest, get the next char otherwise
            if inComment == 0 {
                singleChar = libgogo.GetChar(fd);
                if singleChar == '/' {
                    // we are in a single line comment (until newline is found)
                    inComment = 1;
                } else {
                    if singleChar == '*' {
                        // we are in a multiline comment (until ending is found)
                        inComment = 2;
                    } else {
                        libgogo.ExitError(">> Scanner: Unkown character combination for comments. Exiting.",1);
                    }
                }
            }
        }

        // check whether a multi-line comment is ending
        if singleChar == '*' {
            singleChar = libgogo.GetChar(fd);
            if singleChar == '/' {
                if inComment == 2 {
                    inComment = 0;
                    singleChar = libgogo.GetChar(fd);
                }
            }
        }

        // if character is a newline:
        //  *) if in a singleline comment, exit the comment
        //  *) skip otherwise
        if singleChar == 10 {
            if inComment == 1 {
                inComment = 0;
            }
        }

        // handle everything that is not a space,tab,newline
        if singleChar != ' ' && singleChar != 9 && singleChar != 10 {
            // if not in a comment we have our current valid char
            if inComment == 0 {
                spaceDone = 1;
            }

            // check if GetChar() returned EOF while skipping
            if singleChar == 0 {
                tok.id = TOKEN_EOS;
                spaceDone = 1;
                done = 1;
            }
        }

        // if we are not done until now, get a new character and start another abolishing cycle
        if spaceDone == 0 {
            singleChar=libgogo.GetChar(fd);
        }
    }

    //
    // Actual scanning part starts here
    //

    // Catch identifiers
    // identifier = letter { letter | digit }.
    if (done != 1) && (singleChar >= 'A' && singleChar <= 'Z') || (singleChar >= 'a' && singleChar <= 'z') || singleChar == '_' { // check for letter or _
        tok.id = TOKEN_IDENTIFIER;
        // preceding characters may be letter,_, or a number
        for ; (singleChar >= 'A' && singleChar <= 'Z') || (singleChar >= 'a' && singleChar <= 'z') || singleChar == '_' || (singleChar >= '0' && singleChar <= '9'); singleChar = libgogo.GetChar(fd) {
            tmp_TokAppendStr(tok,singleChar);
        }
        // save the last read character for the next GetNextToken() cycle
        tok.nextChar = singleChar;
        done = 1;
    }

    // string "..."
    if (done != 1) && singleChar == '"' {
        tok.id = TOKEN_STRING;
        for singleChar = libgogo.GetChar(fd); singleChar != '"' &&singleChar > 31 && singleChar < 127;singleChar = libgogo.GetChar(fd) {
            tmp_TokAppendStr(tok,singleChar);
        }
        if singleChar != '"' {
            libgogo.ExitError(">> Scanner: String not closing. Exiting.",1);
        }
        done = 1;
    }

    // Single Quoted Character
    if (done != 1) && singleChar == 39 {
        singleChar = libgogo.GetChar(fd);
        if singleChar != 39 && singleChar > 31 && singleChar < 127 {
            tok.id = TOKEN_INTEGER;
            tok.intValue = libgogo.ToIntFromByte(singleChar);
        } else {
            libgogo.ExitError(">> Scanner: Unknown character. Exiting.",1);
        }
        singleChar = libgogo.GetChar(fd);
        if singleChar != 39 {
            libgogo.ExitError(">> Scanner: Only single characters allowed. Use corresponding integer for special characters. Exiting.",1);
        }
        done = 1;
    }

    // left brace (
    if (done != 1) && singleChar == '(' {
        tok.id = TOKEN_LBRAC;
        done = 1;
    }

    // right brace )
    if (done != 1) && singleChar == ')' {
        tok.id = TOKEN_RBRAC;
        done = 1;
    }

    // left square bracket [
    if (done != 1) && singleChar == '[' {
        tok.id = TOKEN_LSBRAC;
        done = 1;
    }

    // right square bracket ]
    if (done != 1) && singleChar == ']' {
        tok.id = TOKEN_RSBRAC;
        done = 1;
    }

    // integer
    if (done != 1) && singleChar > 47 && singleChar < 58 {
        var numBuf string = "";

        for ; singleChar > 47 && singleChar < 58 ; singleChar = libgogo.GetChar(fd) {
            libgogo.StringAppend(&numBuf, singleChar);
        }

        tok.nextChar = singleChar;
        tok.id = TOKEN_INTEGER;
        tok.intValue = libgogo.StringToInt(numBuf);

        done = 1;
    }

    // Left curly bracket '{'
    if (done != 1) && singleChar == '{' {
        tok.id = TOKEN_LCBRAC;
        done = 1;
    }

    // Right curly bracket '}'
    if (done != 1) && singleChar == '}' {
        tok.id = TOKEN_RCBRAC;
        done = 1;
    }

    // Point '.'
    if (done != 1) && singleChar == '.' {
        tok.id = TOKEN_PT;
        done = 1;
    }

    // Not ('!') or Not Equal ('!=')
    if (done != 1) && singleChar == '!' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '=' {
            tok.id = TOKEN_NOTEQUAL;
        } else {
            tok.id = TOKEN_NOT;
            tok.nextChar = singleChar;
        }
        done = 1;
    }

    // Semicolon ';'
    if (done != 1) && singleChar == ';' {
        tok.id = TOKEN_SEMICOLON;
        done = 1;
    }

    // Colon ','
    if (done != 1) && singleChar == ',' {
        tok.id = TOKEN_COLON;
        done = 1;
    }

    // Assignment '=' or Equals comparison '=='
    if (done != 1) && singleChar == '=' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '=' {
            tok.id = TOKEN_EQUALS;
        } else {
            tok.id = TOKEN_ASSIGN;
            tok.nextChar = singleChar;
        }
        done = 1;
    }

    // AND Relation '&&'
    if (done != 1) && singleChar == '&' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '&' {
            tok.id = TOKEN_REL_AND;
        } else {
            tok.id = TOKEN_OP_ADR;
            tok.nextChar = singleChar;
        }
        done = 1;
    }

    // OR Relation '||'
    if (done != 1) && singleChar == '|' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '|' {
            tok.id = TOKEN_REL_OR;
        } else {
            libgogo.ExitError(">> Scanner: No binary OR (|) supported. Only ||.",1);
        }
        done = 1;
    }

    // Greater and Greater-Than relation
    if (done != 1) && singleChar == '>' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '=' {
            tok.id = TOKEN_REL_GTOE;
        } else {
            tok.id = TOKEN_REL_GT;
            tok.nextChar = singleChar;
        }
        done = 1;
    }

    // Less and Less-Than relation
    if (done != 1) && singleChar == '<' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '=' {
            tok.id = TOKEN_REL_LTOE;
        } else {
            tok.id = TOKEN_REL_LT;
            tok.nextChar = singleChar;
        }
        done = 1;
    }

    if (done != 1) && singleChar == '+' {
        tok.id = TOKEN_ARITH_PLUS;
        done = 1;
    }

    if (done != 1) && singleChar == '-' {
        tok.id = TOKEN_ARITH_MINUS;
        done = 1;
    }

    if (done != 1) && singleChar == '*' {
        tok.id = TOKEN_ARITH_MUL;
        done = 1;
    }

    if (done != 1) && singleChar == '/' {
        tok.id = TOKEN_ARITH_DIV;
        done = 1;
    }

    if done != 1 {

        libgogo.PrintString(">> Scanner: Unkown char '");
        libgogo.PrintChar(singleChar);
        libgogo.PrintString("'. ");
        libgogo.ExitError("Exiting.",1);
    }
}

//
// GetNextToken should be called by the parser. It bascially fetches the next
// token by calling GetNextTokenRaw() and filters the identifiers for known
// keywords.
//
func GetNextToken(fd uint64, tok *Token) {
    GetNextTokenRaw(fd,tok)

    // Convert identifier to keyworded tokens
    if tok.id == TOKEN_IDENTIFIER {
        if libgogo.StringCompare("if",tok.strValue) != 0 {
            tok.id = TOKEN_IF;
       }
        if libgogo.StringCompare("for",tok.strValue) != 0 {
            tok.id = TOKEN_FOR;
        }
        if libgogo.StringCompare("type",tok.strValue) != 0 {
            tok.id = TOKEN_TYPE;
        }
        if libgogo.StringCompare("const",tok.strValue) != 0 {
            tok.id = TOKEN_CONST;
        }
        if libgogo.StringCompare("var",tok.strValue) != 0 {
            tok.id = TOKEN_VAR;
        }
        if libgogo.StringCompare("struct", tok.strValue) != 0 {
            tok.id = TOKEN_STRUCT;
        }
        if libgogo.StringCompare("return", tok.strValue) != 0 {
            tok.id = TOKEN_RETURN;
        }
        if libgogo.StringCompare("func", tok.strValue) != 0 {
            tok.id = TOKEN_FUNC;
        }
        if libgogo.StringCompare("import", tok.strValue) != 0 {
            tok.id = TOKEN_IMPORT;
        }
        if libgogo.StringCompare("package", tok.strValue) != 0 {
            tok.id = TOKEN_PACKAGE;
        }
    }
}

//
// Debugging and temporary functions
//

func debugToken(tok *Token) {
    libgogo.PrintString("---------------------\n");
    libgogo.PrintString("Token Id: ");
    libgogo.PrintNumber(tok.id);
    libgogo.PrintString("\n");
    if tok.id == TOKEN_IDENTIFIER || tok.id == TOKEN_STRING {
        libgogo.PrintString("Stored string: ");
        libgogo.PrintString(tok.strValue);
        libgogo.PrintString("\n");
    }
    if tok.id == TOKEN_INTEGER {
        libgogo.PrintString("Stored integer: ");
        libgogo.PrintNumber(tok.intValue);
        libgogo.PrintString("\n");
    }
}

// Temporary test function
func ScannerTest(fd uint64) {
    var tok Token;

    tok.id = 0;
    tok.nextChar = 0;

    for GetNextToken(fd,&tok); tok.id != TOKEN_EOS; GetNextToken(fd,&tok) {
        debugToken(&tok);
    }
}

func tmp_TokAppendStr(tok *Token, b byte) {
    libgogo.StringAppend(&tok.strValue, b);
}
<commit_msg>scanner.go: Fix StringCompare to match with libgogo<commit_after>// Copyright 2010 The GoGo Authors. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

// This file holds the basic scanning routines that separate a source file
// into the various tokens

package main

import "./libgogo/_obj/libgogo"

// Token struct holding the relevant data of a parsed token.
type Token struct {
    id uint64; // The id. Is one of TOKEN_*
    intValue uint64; // value storing the integer value if the token is TOKEN_INTEGER
    strValue string; // Value storing the token string if the token is TOKEN_STRING or TOKEN_IDENTIFIER
    nextChar byte; // Sometime the next char is already read. It is stored here to be re-assigned in the next GetNextToken() round
    nextToken uint64;
};

var lineCounter uint64 = 0;

func GetNextTokenRaw(fd uint64, tok *Token) {
    var singleChar byte; // Byte holding the last read value
    // Flag indicating whether we are in a comment.
    // 0 for no comment
    // 1 for a single line comment
    // 2 for a multi line comment
    var inComment uint64;
    var done uint64; // Flag indicating whether a cycle (Token) is finsihed
    var spaceDone uint64; // Flag indicating whether an abolishment cycle is finished

    // Initialize variables
    done = 0;
    spaceDone = 0;
    inComment = 0;

    tok.strValue = "";

    // If the previous cycle had to read the next char (and stored it), it is
    // now used as first read
    if tok.nextChar == 0 {
        singleChar = libgogo.GetChar(fd)
    } else {
        singleChar = tok.nextChar;
        tok.nextChar = 0;
    }

    // check if it is a valid read, or an EOF
    if singleChar == 0 {
        tok.id = TOKEN_EOS;
        done = 1;
        spaceDone = 1;
    }

    //
    // Cleaning Tasks
    // The next part strips out spaces, newlines, tabs, and comments
    // Comments can either be single line with double slashes (//) or multiline
    // using C++ syntax /* */
    //
    for ; spaceDone != 1; {

        // check whether a comment is starting
        if singleChar == '/' {
            // if we are in a comment skip the rest, get the next char otherwise
            if inComment == 0 {
                singleChar = libgogo.GetChar(fd);
                if singleChar == '/' {
                    // we are in a single line comment (until newline is found)
                    inComment = 1;
                } else {
                    if singleChar == '*' {
                        // we are in a multiline comment (until ending is found)
                        inComment = 2;
                    } else {
                        libgogo.ExitError(">> Scanner: Unkown character combination for comments. Exiting.",1);
                    }
                }
            }
        }

        // check whether a multi-line comment is ending
        if singleChar == '*' {
            singleChar = libgogo.GetChar(fd);
            if singleChar == '/' {
                if inComment == 2 {
                    inComment = 0;
                    singleChar = libgogo.GetChar(fd);
                }
            }
        }

        // if character is a newline:
        //  *) if in a singleline comment, exit the comment
        //  *) skip otherwise
        if singleChar == 10 {
            lineCounter = lineCounter + 1;
            if inComment == 1 {
                inComment = 0;
            }
        }

        // handle everything that is not a space,tab,newline
        if singleChar != ' ' && singleChar != 9 && singleChar != 10 {
            // if not in a comment we have our current valid char
            if inComment == 0 {
                spaceDone = 1;
            }

            // check if GetChar() returned EOF while skipping
            if singleChar == 0 {
                tok.id = TOKEN_EOS;
                spaceDone = 1;
                done = 1;
            }
        }

        // if we are not done until now, get a new character and start another abolishing cycle
        if spaceDone == 0 {
            singleChar=libgogo.GetChar(fd);
        }
    }

    //
    // Actual scanning part starts here
    //

    // Catch identifiers
    // identifier = letter { letter | digit }.
    if (done != 1) && (singleChar >= 'A' && singleChar <= 'Z') || (singleChar >= 'a' && singleChar <= 'z') || singleChar == '_' { // check for letter or _
        tok.id = TOKEN_IDENTIFIER;
        // preceding characters may be letter,_, or a number
        for ; (singleChar >= 'A' && singleChar <= 'Z') || (singleChar >= 'a' && singleChar <= 'z') || singleChar == '_' || (singleChar >= '0' && singleChar <= '9'); singleChar = libgogo.GetChar(fd) {
            tmp_TokAppendStr(tok,singleChar);
        }
        // save the last read character for the next GetNextToken() cycle
        tok.nextChar = singleChar;
        done = 1;
    }

    // string "..."
    if (done != 1) && singleChar == '"' {
        tok.id = TOKEN_STRING;
        for singleChar = libgogo.GetChar(fd); singleChar != '"' &&singleChar > 31 && singleChar < 127;singleChar = libgogo.GetChar(fd) {
            tmp_TokAppendStr(tok,singleChar);
        }
        if singleChar != '"' {
            libgogo.ExitError(">> Scanner: String not closing. Exiting.",1);
        }
        done = 1;
    }

    // Single Quoted Character
    if (done != 1) && singleChar == 39 {
        singleChar = libgogo.GetChar(fd);
        if singleChar != 39 && singleChar > 31 && singleChar < 127 {
            tok.id = TOKEN_INTEGER;
            tok.intValue = libgogo.ToIntFromByte(singleChar);
        } else {
            libgogo.ExitError(">> Scanner: Unknown character. Exiting.",1);
        }
        singleChar = libgogo.GetChar(fd);
        if singleChar != 39 {
            libgogo.ExitError(">> Scanner: Only single characters allowed. Use corresponding integer for special characters. Exiting.",1);
        }
        done = 1;
    }

    // left brace (
    if (done != 1) && singleChar == '(' {
        tok.id = TOKEN_LBRAC;
        done = 1;
    }

    // right brace )
    if (done != 1) && singleChar == ')' {
        tok.id = TOKEN_RBRAC;
        done = 1;
    }

    // left square bracket [
    if (done != 1) && singleChar == '[' {
        tok.id = TOKEN_LSBRAC;
        done = 1;
    }

    // right square bracket ]
    if (done != 1) && singleChar == ']' {
        tok.id = TOKEN_RSBRAC;
        done = 1;
    }

    // integer
    if (done != 1) && singleChar > 47 && singleChar < 58 {
        var numBuf string = "";

        for ; singleChar > 47 && singleChar < 58 ; singleChar = libgogo.GetChar(fd) {
            libgogo.StringAppend(&numBuf, singleChar);
        }

        tok.nextChar = singleChar;
        tok.id = TOKEN_INTEGER;
        tok.intValue = libgogo.StringToInt(numBuf);

        done = 1;
    }

    // Left curly bracket '{'
    if (done != 1) && singleChar == '{' {
        tok.id = TOKEN_LCBRAC;
        done = 1;
    }

    // Right curly bracket '}'
    if (done != 1) && singleChar == '}' {
        tok.id = TOKEN_RCBRAC;
        done = 1;
    }

    // Point '.'
    if (done != 1) && singleChar == '.' {
        tok.id = TOKEN_PT;
        done = 1;
    }

    // Not ('!') or Not Equal ('!=')
    if (done != 1) && singleChar == '!' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '=' {
            tok.id = TOKEN_NOTEQUAL;
        } else {
            tok.id = TOKEN_NOT;
            tok.nextChar = singleChar;
        }
        done = 1;
    }

    // Semicolon ';'
    if (done != 1) && singleChar == ';' {
        tok.id = TOKEN_SEMICOLON;
        done = 1;
    }

    // Colon ','
    if (done != 1) && singleChar == ',' {
        tok.id = TOKEN_COLON;
        done = 1;
    }

    // Assignment '=' or Equals comparison '=='
    if (done != 1) && singleChar == '=' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '=' {
            tok.id = TOKEN_EQUALS;
        } else {
            tok.id = TOKEN_ASSIGN;
            tok.nextChar = singleChar;
        }
        done = 1;
    }

    // AND Relation '&&'
    if (done != 1) && singleChar == '&' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '&' {
            tok.id = TOKEN_REL_AND;
        } else {
            tok.id = TOKEN_OP_ADR;
            tok.nextChar = singleChar;
        }
        done = 1;
    }

    // OR Relation '||'
    if (done != 1) && singleChar == '|' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '|' {
            tok.id = TOKEN_REL_OR;
        } else {
            libgogo.ExitError(">> Scanner: No binary OR (|) supported. Only ||.",1);
        }
        done = 1;
    }

    // Greater and Greater-Than relation
    if (done != 1) && singleChar == '>' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '=' {
            tok.id = TOKEN_REL_GTOE;
        } else {
            tok.id = TOKEN_REL_GT;
            tok.nextChar = singleChar;
        }
        done = 1;
    }

    // Less and Less-Than relation
    if (done != 1) && singleChar == '<' {
        singleChar = libgogo.GetChar(fd);
        if singleChar == '=' {
            tok.id = TOKEN_REL_LTOE;
        } else {
            tok.id = TOKEN_REL_LT;
            tok.nextChar = singleChar;
        }
        done = 1;
    }

    if (done != 1) && singleChar == '+' {
        tok.id = TOKEN_ARITH_PLUS;
        done = 1;
    }

    if (done != 1) && singleChar == '-' {
        tok.id = TOKEN_ARITH_MINUS;
        done = 1;
    }

    if (done != 1) && singleChar == '*' {
        tok.id = TOKEN_ARITH_MUL;
        done = 1;
    }

    if (done != 1) && singleChar == '/' {
        tok.id = TOKEN_ARITH_DIV;
        done = 1;
    }

    if done != 1 {

        libgogo.PrintString(">> Scanner: Unkown char '");
        libgogo.PrintChar(singleChar);
        libgogo.PrintString("'. ");
        libgogo.ExitError("Exiting.",1);
    }
}

//
// GetNextToken should be called by the parser. It bascially fetches the next
// token by calling GetNextTokenRaw() and filters the identifiers for known
// keywords.
//
func GetNextToken(fd uint64, tok *Token) {
    GetNextTokenRaw(fd, tok);

    // Convert identifier to keyworded tokens
    if tok.id == TOKEN_IDENTIFIER {
        if libgogo.StringCompare("if",tok.strValue) == 0 {
            tok.id = TOKEN_IF;
       }
        if libgogo.StringCompare("for",tok.strValue) == 0 {
            tok.id = TOKEN_FOR;
        }
        if libgogo.StringCompare("type",tok.strValue) == 0 {
            tok.id = TOKEN_TYPE;
        }
        if libgogo.StringCompare("const",tok.strValue) == 0 {
            tok.id = TOKEN_CONST;
        }
        if libgogo.StringCompare("var",tok.strValue) == 0 {
            tok.id = TOKEN_VAR;
        }
        if libgogo.StringCompare("struct", tok.strValue) == 0 {
            tok.id = TOKEN_STRUCT;
        }
        if libgogo.StringCompare("return", tok.strValue) == 0 {
            tok.id = TOKEN_RETURN;
        }
        if libgogo.StringCompare("func", tok.strValue) == 0 {
            tok.id = TOKEN_FUNC;
        }
        if libgogo.StringCompare("import", tok.strValue) == 0 {
            tok.id = TOKEN_IMPORT;
        }
        if libgogo.StringCompare("package", tok.strValue) == 0 {
            tok.id = TOKEN_PACKAGE;
        }
    }

    tok.nextToken = 0;
}

//
// Debugging and temporary functions
//

func debugToken(tok *Token) {
    libgogo.PrintString("---------------------\n");
    libgogo.PrintString("Token Id: ");
    libgogo.PrintNumber(tok.id);
    libgogo.PrintString("\n");
    if tok.id == TOKEN_IDENTIFIER || tok.id == TOKEN_STRING {
        libgogo.PrintString("Stored string: ");
        libgogo.PrintString(tok.strValue);
        libgogo.PrintString("\n");
    }
    if tok.id == TOKEN_INTEGER {
        libgogo.PrintString("Stored integer: ");
        libgogo.PrintNumber(tok.intValue);
        libgogo.PrintString("\n");
    }
}

// Temporary test function
func ScannerTest(fd uint64) {
    var tok Token;

    tok.id = 0;
    tok.nextChar = 0;

    for GetNextToken(fd,&tok); tok.id != TOKEN_EOS; GetNextToken(fd,&tok) {
        debugToken(&tok);
    }
}

func tmp_TokAppendStr(tok *Token, b byte) {
    libgogo.StringAppend(&tok.strValue, b);
}
<|endoftext|>


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

package src

import (
	"errors"
	"fmt"
	"io"
)

const bufsize = 8192

const (
	scanIllegalToken = iota
	scanBeginObject
	scanEndObject
	scanBeginArray
	scanEndArray
	scanIntLit
	scanStringLit
	scanBoolLit
	scanNumberLit
	scanNullVal
	scanComma
)

type scanner struct {
	r   io.Reader
	err error

	row int64
	col int64

	buf []byte
	pos int // position inside the buffer; must be -1 by default

	eof bool // true when the reader has been reached
}

func newScanner(r io.Reader) *scanner {
	return &scanner{r: r, buf: make([]byte, bufsize), pos: -1}
}

func (scan *scanner) nextKey() (string, error) {
	if err := scan.ignoreWhitespaces(); err != nil {
		if err == io.EOF {
			return "", errors.New("expected key, found EOF")
		}
		return "", err
	}

	if c, err := scan.read(); err != nil {
		return "", err
	} else if c != '"' {
		return "", fmt.Errorf("expected '\"', found '%c'", c)
	}

	var key string
	origPos := scan.pos
	for {
		c, err := scan.read()
		if err == io.EOF {
			return "", errors.New("expected key, found EOF")
		} else if err != nil {
			return "", err
		}

		// TODO: check character

		if c == '"' {
			break
		}

		if scan.pos >= len(scan.buf)-1 {
			key += string(scan.buf[origPos:])
			origPos = 0
		}
	}

	key += string(scan.buf[origPos:scan.pos])

	if err := scan.ignoreWhitespaces(); err != nil {
		if err == io.EOF {
			return "", errors.New("expected ':', found EOF")
		}
		return "", err
	}

	if c, err := scan.read(); err == io.EOF {
		return "", errors.New("expected ':', found EOF")
	} else if err != nil {
		return "", err
	} else if c != ':' {
		return "", fmt.Errorf("expected ':', found '%c'", c)
	}

	return key, nil
}

func (scan *scanner) nextValue() (val []byte, tok int, err error) {
	if err := scan.ignoreWhitespaces(); err != nil {
		if err == io.EOF {
			return nil, 0, errors.New("expected value, found EOF")
		}
		return nil, 0, err
	}

	c, err := scan.read()
	if err != nil {
		if err == io.EOF {
			return nil, 0, errors.New("expected value, found EOF")
		}
		return nil, 0, err
	}

	switch {
	case isDigit(c):
		// TODO: read digit
	case c == 't' || c == 'f':
		// TODO: read bool
	case c == '{': // beginning of an object literal
		return nil, scanBeginObject, nil
	case c == '}': // ending of an object literal
		return nil, scanEndObject, nil
	case c == '[': // beginning of an array literal
		return nil, scanBeginArray, nil
	case c == ']': // ending of an array literal
		return nil, scanEndArray, nil
	case c == '"': // beginning or ending of a string literal
		// TODO: read string
	case c == 'n': // null
		// TODO: read null
	case c == ',':
		return nil, scanComma, nil
	}
	return nil, scanIllegalToken, fmt.Errorf("illegal value '%c'", c)
}

// peek reads the next value without consuming it.
func (scan *scanner) peek() (byte, error) {
	c, err := scan.read()
	if err != nil {
		return 0, err
	}
	scan.back()
	return c, nil
}

func (scan *scanner) read() (byte, error) {
	if scan.pos == -1 || (scan.pos > len(scan.buf)-1 && !scan.eof) {
		n, err := scan.r.Read(scan.buf)
		if err != nil {
			if err != io.EOF {
				return 0, err
			}
			scan.eof = true
		}
		if n < bufsize {
			scan.buf = scan.buf[:n]
		}
		scan.pos = 0
	}
	if scan.pos >= len(scan.buf) {
		return 0, io.EOF
	}
	b := scan.buf[scan.pos]
	scan.pos++
	return b, nil
}

func (scan *scanner) back() {
	if scan.pos > 0 {
		scan.pos--
	}
}

// ignoreWhitespaces consumes all whitespaces.
func (scan *scanner) ignoreWhitespaces() error {
	var c byte
	var err error
	// XXX: refactor when the tests pass
	for {
		if c, err = scan.read(); err != nil || !isWhitespace(c) {
			break
		}
		// Nothing, we just skip whitespaces.
	}

	// Since a non-whitespace has been read, we have to put it back to the
	// buffer so that it can be read again.
	scan.back()

	if err != io.EOF {
		return err
	}
	return nil
}

func isWhitespace(c byte) bool {
	return c == ' ' || c == '\n' || c == '\t' || c == '\r'
}

func isDigit(c byte) bool {
	return c >= '0' && c <= '9'
}

// TODO: implement
func isUnicodeChar(c byte) bool {
	return false
}
<commit_msg>src: add string reading support in the scanner<commit_after>// Copyright 2014-2015 The project AUTHORS. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package src

import (
	"errors"
	"fmt"
	"io"
)

const bufsize = 8192

const (
	scanIllegalToken = iota
	scanBeginObject
	scanEndObject
	scanBeginArray
	scanEndArray
	scanIntLit
	scanStringLit
	scanBoolLit
	scanNumberLit
	scanNullVal
	scanComma
)

type scanner struct {
	r   io.Reader
	err error

	row int64
	col int64

	buf []byte
	pos int // position inside the buffer; must be -1 by default

	eof bool // true when the reader has been reached
}

func newScanner(r io.Reader) *scanner {
	return &scanner{r: r, buf: make([]byte, bufsize), pos: -1}
}

func (scan *scanner) nextKey() (string, error) {
	if err := scan.ignoreWhitespaces(); err != nil {
		if err == io.EOF {
			return "", errors.New("expected key, found EOF")
		}
		return "", err
	}

	if c, err := scan.read(); err != nil {
		return "", err
	} else if c != '"' {
		return "", fmt.Errorf("expected '\"', found '%c'", c)
	}

	var key string
	origPos := scan.pos
	for {
		c, err := scan.read()
		if err == io.EOF {
			return "", errors.New("expected key, found EOF")
		} else if err != nil {
			return "", err
		}

		// TODO: check character

		if c == '"' {
			break
		}

		if scan.pos >= len(scan.buf)-1 {
			key += string(scan.buf[origPos:])
			origPos = 0
		}
	}

	key += string(scan.buf[origPos:scan.pos])

	if err := scan.ignoreWhitespaces(); err != nil {
		if err == io.EOF {
			return "", errors.New("expected ':', found EOF")
		}
		return "", err
	}

	if c, err := scan.read(); err == io.EOF {
		return "", errors.New("expected ':', found EOF")
	} else if err != nil {
		return "", err
	} else if c != ':' {
		return "", fmt.Errorf("expected ':', found '%c'", c)
	}

	return key, nil
}

func (scan *scanner) nextValue() (val []byte, tok int, err error) {
	if err := scan.ignoreWhitespaces(); err != nil {
		if err == io.EOF {
			return nil, 0, errors.New("expected value, found EOF")
		}
		return nil, 0, err
	}

	c, err := scan.read()
	if err != nil {
		if err == io.EOF {
			return nil, 0, errors.New("expected value, found EOF")
		}
		return nil, 0, err
	}

	switch {
	case isDigit(c):
		// TODO: read digit
	case c == 't' || c == 'f':
		// TODO: read bool
	case c == '{': // beginning of an object literal
		return nil, scanBeginObject, nil
	case c == '}': // ending of an object literal
		return nil, scanEndObject, nil
	case c == '[': // beginning of an array literal
		return nil, scanBeginArray, nil
	case c == ']': // ending of an array literal
		return nil, scanEndArray, nil
	case c == '"': // beginning or ending of a string literal
		str, err := scan.readString()
		return str, scanStringLit, err
	case c == 'n': // null
		// TODO: read null
	case c == ',':
		return nil, scanComma, nil
	}
	return nil, scanIllegalToken, fmt.Errorf("illegal value '%c'", c)
}

// peek reads the next value without consuming it.
func (scan *scanner) peek() (byte, error) {
	c, err := scan.read()
	if err != nil {
		return 0, err
	}
	scan.back()
	return c, nil
}

func (scan *scanner) read() (byte, error) {
	if scan.pos == -1 || (scan.pos > len(scan.buf)-1 && !scan.eof) {
		n, err := scan.r.Read(scan.buf)
		if err != nil {
			if err != io.EOF {
				return 0, err
			}
			scan.eof = true
		}
		if n < bufsize {
			scan.buf = scan.buf[:n]
		}
		scan.pos = 0
	}
	if scan.pos >= len(scan.buf) {
		return 0, io.EOF
	}
	b := scan.buf[scan.pos]
	scan.pos++
	return b, nil
}

// readString reads a string literal.
func (scan *scanner) readString() ([]byte, error) {
	// length of the read string
	// XXX: overflow?
	var strLen int

	buf := make([]byte, bufsize)
	for {
		c, err := scan.read()
		if err != nil {
			return nil, err
		}
		if c == '"' {
			if strLen == 0 {
				// empty string
				return []byte{}, nil
			}
			if buf[strLen-1] != '\\' {
				break
			}
		}
		if strLen == len(buf) {
			buf = append(buf, c)
		} else {
			buf[strLen] = c
		}
		strLen++
	}

	return buf[:strLen], nil
}

func (scan *scanner) back() {
	if scan.pos > 0 {
		scan.pos--
	}
}

// ignoreWhitespaces consumes all whitespaces.
func (scan *scanner) ignoreWhitespaces() error {
	var c byte
	var err error
	// XXX: refactor when the tests pass
	for {
		if c, err = scan.read(); err != nil || !isWhitespace(c) {
			break
		}
		// Nothing, we just skip whitespaces.
	}

	// Since a non-whitespace has been read, we have to put it back to the
	// buffer so that it can be read again.
	scan.back()

	if err != io.EOF {
		return err
	}
	return nil
}

func isWhitespace(c byte) bool {
	return c == ' ' || c == '\n' || c == '\t' || c == '\r'
}

func isDigit(c byte) bool {
	return c >= '0' && c <= '9'
}

// TODO: implement
func isUnicodeChar(c byte) bool {
	return false
}
<|endoftext|>


<commit_before>// Copyright 2015 The rkt Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build host coreos src kvm

package main

import (
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/rkt/rkt/tests/testutils"
)

const delta = 4 * time.Second
const precision = 2 * time.Second

// compareTime checks if a and b are roughly equal
func compareTime(a time.Time, b time.Time) bool {
	diff := a.Sub(b)
	if diff < 0 {
		diff = -diff
	}
	return diff < precision
}

func TestRktList(t *testing.T) {
	const imgName = "rkt-list-test"

	image := patchTestACI(fmt.Sprintf("%s.aci", imgName), fmt.Sprintf("--name=%s", imgName))
	defer os.Remove(image)

	imageHash := getHashOrPanic(image)
	imgID := ImageID{image, imageHash}

	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	// Prepare image
	cmd := fmt.Sprintf("%s --insecure-options=image prepare %s", ctx.Cmd(), imgID.path)
	podUuid := runRktAndGetUUID(t, cmd)

	// Get hash
	imageID := fmt.Sprintf("sha512-%s", imgID.hash[:12])

	// Define tests
	tests := []struct {
		cmd           string
		shouldSucceed bool
		expect        string
	}{
		// Test that pod UUID is in output
		{
			"list --full",
			true,
			podUuid,
		},
		// Test that image name is in output
		{
			"list",
			true,
			imgName,
		},
		// Test that imageID is in output
		{
			"list --full",
			true,
			imageID,
		},
		// Remove the image
		{
			fmt.Sprintf("image rm %s", imageID),
			true,
			"successfully removed",
		},
		// Name should still show up in rkt list
		{
			"list",
			true,
			imgName,
		},
		// Test that imageID is still in output
		{
			"list --full",
			true,
			imageID,
		},
	}

	// Run tests
	for i, tt := range tests {
		runCmd := fmt.Sprintf("%s %s", ctx.Cmd(), tt.cmd)
		t.Logf("Running test #%d, %s", i, runCmd)
		runRktAndCheckOutput(t, runCmd, tt.expect, !tt.shouldSucceed)
	}
}

func getCreationStartTime(t *testing.T, ctx *testutils.RktRunCtx, imageID string) (creation time.Time, start time.Time) {
	// Run rkt list --full
	rktCmd := fmt.Sprintf("%s list --full", ctx.Cmd())
	child := spawnOrFail(t, rktCmd)
	child.Wait()

	// Get creation time
	match := fmt.Sprintf(".*%s\t.*\t(.*)\t(.*)\t", imageID)
	result, out, err := expectRegexWithOutput(child, match)
	if err != nil {
		t.Fatalf("%q regex not found, Error: %v\nOutput: %v", match, err, out)
	}
	tmStr := strings.TrimSpace(result[1])
	creation, err = time.Parse(defaultTimeLayout, tmStr)
	if err != nil {
		t.Fatalf("Error parsing creation time: %q", err)
	}

	tmStr = strings.TrimSpace(result[2])
	start, err = time.Parse(defaultTimeLayout, tmStr)
	if err != nil {
		t.Fatalf("Error parsing start time: %q", err)
	}

	return creation, start
}

func TestRktListCreatedStarted(t *testing.T) {
	const imgName = "rkt-list-creation-start-time-test"

	image := patchTestACI(fmt.Sprintf("%s.aci", imgName), fmt.Sprintf("--exec=/inspect --exit-code=0"))
	defer os.Remove(image)

	imageHash := getHashOrPanic(image)
	imgID := ImageID{image, imageHash}

	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	// Prepare image
	cmd := fmt.Sprintf("%s --insecure-options=image prepare %s", ctx.Cmd(), imgID.path)
	podUuid := runRktAndGetUUID(t, cmd)

	// t0: prepare
	expectPrepare := time.Now()

	// Get hash
	imageID := fmt.Sprintf("sha512-%s", imgID.hash[:12])

	tmpDir := mustTempDir(imgName)
	defer os.RemoveAll(tmpDir)

	time.Sleep(delta)

	// Run image
	cmd = fmt.Sprintf("%s run-prepared %s", ctx.Cmd(), podUuid)
	rktChild := spawnOrFail(t, cmd)

	// t1: run
	expectRun := time.Now()

	waitOrFail(t, rktChild, 0)

	creation, start := getCreationStartTime(t, ctx, imageID)
	if !compareTime(expectPrepare, creation) {
		t.Fatalf("rkt list returned an incorrect creation time. Got: %q Expect: %q", creation, expectPrepare)
	}
	if !compareTime(expectRun, start) {
		t.Fatalf("rkt list returned an incorrect start time. Got: %q Expect: %q", start, expectRun)
	}
}
<commit_msg>tests: check ip of non running pod for list command<commit_after>// Copyright 2015 The rkt Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build host coreos src kvm

package main

import (
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/rkt/rkt/tests/testutils"
)

const delta = 4 * time.Second
const precision = 2 * time.Second

// compareTime checks if a and b are roughly equal
func compareTime(a time.Time, b time.Time) bool {
	diff := a.Sub(b)
	if diff < 0 {
		diff = -diff
	}
	return diff < precision
}

func TestRktList(t *testing.T) {
	const imgName = "rkt-list-test"

	image := patchTestACI(fmt.Sprintf("%s.aci", imgName), fmt.Sprintf("--name=%s", imgName))
	defer os.Remove(image)

	imageHash := getHashOrPanic(image)
	imgID := ImageID{image, imageHash}

	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	// Prepare image
	cmd := fmt.Sprintf("%s --insecure-options=image prepare %s", ctx.Cmd(), imgID.path)
	podUuid := runRktAndGetUUID(t, cmd)

	// Get hash
	imageID := fmt.Sprintf("sha512-%s", imgID.hash[:12])

	cmd = fmt.Sprintf("%s --insecure-options=image run-prepared %s", ctx.Cmd(), podUuid)
	nobodyUid, _ := testutils.GetUnprivilegedUidGid()
	_, status := runRkt(t, cmd, nobodyUid, 0)
	if status != 0 {
		panic(fmt.Errorf("expected exit status code 0, got %d", status))
	}
	podInfo := getPodInfo(t, ctx, podUuid)
	ipv4 := podInfo.networks["default"].ipv4
	// Define tests
	tests := []struct {
		cmd           string
		shouldSucceed bool
		expect        string
	}{
		// Test that pod UUID is in output
		{
			"list --full",
			true,
			podUuid,
		},
		// Test that image name is in output
		{
			"list",
			true,
			imgName,
		},
		// Test that imageID is in output
		{
			"list --full",
			true,
			imageID,
		},
		// Remove the image
		{
			fmt.Sprintf("image rm %s", imageID),
			true,
			"successfully removed",
		},
		// Name should still show up in rkt list
		{
			"list",
			true,
			imgName,
		},
		// Test that imageID is still in output
		{
			"list --full",
			true,
			imageID,
		},
		//	Test that ip is in json output
		{
			"list --format=json",
			true,
			ipv4,
		},
	}

	// Run tests
	for i, tt := range tests {
		runCmd := fmt.Sprintf("%s %s", ctx.Cmd(), tt.cmd)
		t.Logf("Running test #%d, %s", i, runCmd)
		runRktAndCheckOutput(t, runCmd, tt.expect, !tt.shouldSucceed)
	}
}

func getCreationStartTime(t *testing.T, ctx *testutils.RktRunCtx, imageID string) (creation time.Time, start time.Time) {
	// Run rkt list --full
	rktCmd := fmt.Sprintf("%s list --full", ctx.Cmd())
	child := spawnOrFail(t, rktCmd)
	child.Wait()

	// Get creation time
	match := fmt.Sprintf(".*%s\t.*\t(.*)\t(.*)\t", imageID)
	result, out, err := expectRegexWithOutput(child, match)
	if err != nil {
		t.Fatalf("%q regex not found, Error: %v\nOutput: %v", match, err, out)
	}
	tmStr := strings.TrimSpace(result[1])
	creation, err = time.Parse(defaultTimeLayout, tmStr)
	if err != nil {
		t.Fatalf("Error parsing creation time: %q", err)
	}

	tmStr = strings.TrimSpace(result[2])
	start, err = time.Parse(defaultTimeLayout, tmStr)
	if err != nil {
		t.Fatalf("Error parsing start time: %q", err)
	}

	return creation, start
}

func TestRktListCreatedStarted(t *testing.T) {
	const imgName = "rkt-list-creation-start-time-test"

	image := patchTestACI(fmt.Sprintf("%s.aci", imgName), fmt.Sprintf("--exec=/inspect --exit-code=0"))
	defer os.Remove(image)

	imageHash := getHashOrPanic(image)
	imgID := ImageID{image, imageHash}

	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	// Prepare image
	cmd := fmt.Sprintf("%s --insecure-options=image prepare %s", ctx.Cmd(), imgID.path)
	podUuid := runRktAndGetUUID(t, cmd)

	// t0: prepare
	expectPrepare := time.Now()

	// Get hash
	imageID := fmt.Sprintf("sha512-%s", imgID.hash[:12])

	tmpDir := mustTempDir(imgName)
	defer os.RemoveAll(tmpDir)

	time.Sleep(delta)

	// Run image
	cmd = fmt.Sprintf("%s run-prepared %s", ctx.Cmd(), podUuid)
	rktChild := spawnOrFail(t, cmd)

	// t1: run
	expectRun := time.Now()

	waitOrFail(t, rktChild, 0)

	creation, start := getCreationStartTime(t, ctx, imageID)
	if !compareTime(expectPrepare, creation) {
		t.Fatalf("rkt list returned an incorrect creation time. Got: %q Expect: %q", creation, expectPrepare)
	}
	if !compareTime(expectRun, start) {
		t.Fatalf("rkt list returned an incorrect start time. Got: %q Expect: %q", start, expectRun)
	}
}
<|endoftext|>


<commit_before>// Package alice provides a convenient way to chain http handlers.
package alice

import "net/http"

// A constructor for a piece of middleware.
// Some middleware use this constructor out of the box,
// so in most cases you can just pass somepackage.New
type Constructor func(http.Handler) http.Handler

// Chain acts as a list of http.Handler constructors.
// Chain is effectively immutable:
// once created, it will always hold
// the same set of constructors in the same order.
type Chain struct {
	constructors []Constructor
}

// New creates a new chain,
// memorizing the given list of middleware constructors.
// New serves no other function,
// constructors are only called upon a call to Then().
func New(constructors ...Constructor) Chain {
	c := Chain{}
	c.constructors = append(c.constructors, constructors...)

	return c
}

// Then chains the middleware and returns the final http.Handler.
//     New(m1, m2, m3).Then(h)
// is equivalent to:
//     m1(m2(m3(h)))
// When the request comes in, it will be passed to m1, then m2, then m3
// and finally, the given handler
// (assuming every middleware calls the following one).
//
// A chain can be safely reused by calling Then() several times.
//     stdStack := alice.New(ratelimitHandler, csrfHandler)
//     indexPipe = stdStack.Then(indexHandler)
//     authPipe = stdStack.Then(authHandler)
// Note that constructors are called on every call to Then()
// and thus several instances of the same middleware will be created
// when a chain is reused in this way.
// For proper middleware, this should cause no problems.
//
// Then() treats nil as http.DefaultServeMux.
func (c Chain) Then(h http.Handler) http.Handler {
	var final http.Handler
	if h != nil {
		final = h
	} else {
		final = http.DefaultServeMux
	}

	for i := len(c.constructors) - 1; i >= 0; i-- {
		final = c.constructors[i](final)
	}

	return final
}

// ThenFunc works identically to Then, but takes
// a HandlerFunc instead of a Handler.
//
// The following two statements are equivalent:
//     c.Then(http.HandlerFunc(fn))
//     c.ThenFunc(fn)
//
// ThenFunc provides all the guarantees of Then.
func (c Chain) ThenFunc(fn http.HandlerFunc) http.Handler {
	if fn == nil {
		return c.Then(nil)
	}
	return c.Then(http.HandlerFunc(fn))
}

// Append extends a chain, adding the specified constructors
// as the last ones in the request flow.
//
// Append returns a new chain, leaving the original one untouched.
//
//     stdChain := alice.New(m1, m2)
//     extChain := stdChain.Append(m3, m4)
//     // requests in stdChain go m1 -> m2
//     // requests in extChain go m1 -> m2 -> m3 -> m4
func (c Chain) Append(constructors ...Constructor) Chain {
	newCons := make([]Constructor, len(c.constructors)+len(constructors))
	copy(newCons, c.constructors)
	copy(newCons[len(c.constructors):], constructors)

	newChain := New(newCons...)
	return newChain
}

// Extend extends a chain by adding the specified chain
// as the last one in the request flow.
//
// Extend returns a new chain, leaving the original one untouched.
//
//     stdChain := alice.New(m1, m2)
//     ext1Chain := alice.New(m3, m4)
//     ext2Chain := stdChain.Extend(ext1Chain)
//     // requests in stdChain go  m1 -> m2
//     // requests in ext1Chain go m3 -> m4
//     // requests in ext2Chain go m1 -> m2 -> m3 -> m4
//
// Another example:
//  aHtmlAfterNosurf := alice.New(m2)
// 	aHtml := alice.New(m1, func(h http.Handler) http.Handler {
// 		csrf := nosurf.New(h)
// 		csrf.SetFailureHandler(aHtmlAfterNosurf.ThenFunc(csrfFail))
// 		return csrf
// 	}).Extend(aHtmlAfterNosurf)
//		// requests to aHtml hitting nosurfs success handler go m1 -> nosurf -> m2 -> target-handler
//		// requests to aHtml hitting nosurfs failure handler go m1 -> nosurf -> m2 -> csrfFail
func (c Chain) Extend(chain Chain) Chain {
	return c.Append(chain.constructors...)
}
<commit_msg>refactored New<commit_after>// Package alice provides a convenient way to chain http handlers.
package alice

import "net/http"

// A constructor for a piece of middleware.
// Some middleware use this constructor out of the box,
// so in most cases you can just pass somepackage.New
type Constructor func(http.Handler) http.Handler

// Chain acts as a list of http.Handler constructors.
// Chain is effectively immutable:
// once created, it will always hold
// the same set of constructors in the same order.
type Chain struct {
	constructors []Constructor
}

// New creates a new chain,
// memorizing the given list of middleware constructors.
// New serves no other function,
// constructors are only called upon a call to Then().
func New(constructors ...Constructor) Chain {
	return Chain{append(([]Constructor)(nil), constructors...)}
}

// Then chains the middleware and returns the final http.Handler.
//     New(m1, m2, m3).Then(h)
// is equivalent to:
//     m1(m2(m3(h)))
// When the request comes in, it will be passed to m1, then m2, then m3
// and finally, the given handler
// (assuming every middleware calls the following one).
//
// A chain can be safely reused by calling Then() several times.
//     stdStack := alice.New(ratelimitHandler, csrfHandler)
//     indexPipe = stdStack.Then(indexHandler)
//     authPipe = stdStack.Then(authHandler)
// Note that constructors are called on every call to Then()
// and thus several instances of the same middleware will be created
// when a chain is reused in this way.
// For proper middleware, this should cause no problems.
//
// Then() treats nil as http.DefaultServeMux.
func (c Chain) Then(h http.Handler) http.Handler {
	var final http.Handler
	if h != nil {
		final = h
	} else {
		final = http.DefaultServeMux
	}

	for i := len(c.constructors) - 1; i >= 0; i-- {
		final = c.constructors[i](final)
	}

	return final
}

// ThenFunc works identically to Then, but takes
// a HandlerFunc instead of a Handler.
//
// The following two statements are equivalent:
//     c.Then(http.HandlerFunc(fn))
//     c.ThenFunc(fn)
//
// ThenFunc provides all the guarantees of Then.
func (c Chain) ThenFunc(fn http.HandlerFunc) http.Handler {
	if fn == nil {
		return c.Then(nil)
	}
	return c.Then(http.HandlerFunc(fn))
}

// Append extends a chain, adding the specified constructors
// as the last ones in the request flow.
//
// Append returns a new chain, leaving the original one untouched.
//
//     stdChain := alice.New(m1, m2)
//     extChain := stdChain.Append(m3, m4)
//     // requests in stdChain go m1 -> m2
//     // requests in extChain go m1 -> m2 -> m3 -> m4
func (c Chain) Append(constructors ...Constructor) Chain {
	newCons := make([]Constructor, len(c.constructors)+len(constructors))
	copy(newCons, c.constructors)
	copy(newCons[len(c.constructors):], constructors)

	newChain := New(newCons...)
	return newChain
}

// Extend extends a chain by adding the specified chain
// as the last one in the request flow.
//
// Extend returns a new chain, leaving the original one untouched.
//
//     stdChain := alice.New(m1, m2)
//     ext1Chain := alice.New(m3, m4)
//     ext2Chain := stdChain.Extend(ext1Chain)
//     // requests in stdChain go  m1 -> m2
//     // requests in ext1Chain go m3 -> m4
//     // requests in ext2Chain go m1 -> m2 -> m3 -> m4
//
// Another example:
//  aHtmlAfterNosurf := alice.New(m2)
// 	aHtml := alice.New(m1, func(h http.Handler) http.Handler {
// 		csrf := nosurf.New(h)
// 		csrf.SetFailureHandler(aHtmlAfterNosurf.ThenFunc(csrfFail))
// 		return csrf
// 	}).Extend(aHtmlAfterNosurf)
//		// requests to aHtml hitting nosurfs success handler go m1 -> nosurf -> m2 -> target-handler
//		// requests to aHtml hitting nosurfs failure handler go m1 -> nosurf -> m2 -> csrfFail
func (c Chain) Extend(chain Chain) Chain {
	return c.Append(chain.constructors...)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"github.com/fatih/color"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
)

/// Chasm Types ///

// CloudStore represents an external cloud storage service that is compatible
// with Chasm
type CloudStore interface {
	Upload(share Share)
	Delete(sid ShareID)

	//Restore downloads shares to local restore path
	Restore() string

	Description() string

	Clean()
}

// ChasmPref represents user/application preferences
type ChasmPref struct {
	root string

	// the cloud services sharing across
	FolderStores []FolderStore `json:"folder_stores"`

	// the cloud services sharing across
	GDriveStores []GDriveStore `json:"gdrive_stores"`

	// the cloud services sharing across
	DropboxStores []DropboxStore `json:"dropbox_stores"`

	// maps files to their shareId
	FileMap map[string]ShareID `json:"files"`
}

// RegisteredServices counts all services
func (p ChasmPref) RegisteredServices() int {
	return len(p.FolderStores) + len(p.GDriveStores) + len(p.DropboxStores)
}

// NeedSetup checks if there are enough services to run
func (p ChasmPref) NeedSetup() bool {
	return p.RegisteredServices() < 2
}

// AllCloudStores combines all the cloud stores
func (p ChasmPref) AllCloudStores() []CloudStore {

	// adjust length for new store types
	cloudStores := make([]CloudStore, p.RegisteredServices())

	// all other cloud stores go here
	ind := 0
	for i, fs := range p.FolderStores {
		cloudStores[i] = CloudStore(fs)
		ind += 1
	}

	for j, gds := range p.GDriveStores {
		cloudStores[j+ind] = CloudStore(gds)
		ind += 1
	}

	for k, dbs := range p.DropboxStores {
		cloudStores[k+ind] = CloudStore(dbs)
		ind += 1
	}

	return cloudStores
}

// Save saves the chasm preferences
func (p ChasmPref) Save() {
	chasmFilePath := p.root + string(filepath.Separator) + chasmPrefFile
	chasmFileBytes, err := json.MarshalIndent(preferences, "", "    ")
	check(err)

	ioutil.WriteFile(chasmFilePath, chasmFileBytes, 0660)
}

/// Chasm Functions ///

var preferences ChasmPref

const chasmPrefFile = ".chasm"
const chasmIgnoreFile = ".chasmignore"

// CreateOrLoadChasmDir creates the root *chasm* folder on the system
// if it does not exist or finds an existing directory
// returns if true if newly created
func CreateOrLoadChasmDir(root string) {
	os.MkdirAll(root, 0777)

	chasmFilePath := path.Join(root, chasmPrefFile)
	chasmFileBytes, err := ioutil.ReadFile(chasmFilePath)
	if err != nil {
		color.Green("Creating new .chasm secure folder")
		preferences.FileMap = make(map[string]ShareID)
		preferences.FileMap[chasmFilePath] = ShareID(chasmPrefFile)
	} else {
		json.Unmarshal(chasmFileBytes, &preferences)
	}

	chasmIgnorePath := path.Join(root, chasmIgnoreFile)
	_, err = ioutil.ReadFile(chasmIgnorePath)
	if err != nil {
		preferences.FileMap[chasmIgnorePath] = ShareID(chasmIgnoreFile)
		// add *.DS_Store to ignore file by default
		errWrite := ioutil.WriteFile(chasmIgnorePath, []byte(".DS_Store\n"), 0777)
		if errWrite != nil {
			color.Red("Error: could not write to %s: %s", chasmFilePath, errWrite)
		}
	}

	preferences.root = root
	preferences.Save()
}

// IsValidPath checks if a file path is vaild, i.e. it doesn't match any patterns
// in the .chasmignore file
func IsValidPath(filePath string) bool {
	base := filepath.Base(filePath)
	chasmIgnorePath := path.Join(preferences.root, chasmIgnoreFile)
	chasmIgnore, err := os.Open(chasmIgnorePath)
	if err != nil {
		return true
	}

	scanner := bufio.NewScanner(chasmIgnore)
	for scanner.Scan() {
		pattern := scanner.Text()
		// if the file matches anything in .chasmignore, return false
		ok, err := filepath.Match(pattern, base)
		if ok {
			return false
		}
		if err != nil {
			fmt.Println(err)
		}
	}

	// check for errors
	if err := scanner.Err(); err != nil {
		fmt.Println(err)
	}

	return true
}

// AddFile secret shares the file, and uploads each share to corresponding services
// if the file exists already, we delete the remote share first by its shareId
func AddFile(filePath string) {
	if !IsValidPath(filePath) {
		color.Red("Path %s is in .chasmignore. No actions will be performed.", filePath)
		return
	}
	file, _ := os.Open(filePath)
	fi, err := file.Stat()
	if err != nil {
		color.Red("Cannot get file info: %s", err)
		return
	}
	switch mode := fi.Mode(); {
	case mode.IsDir():
		files, _ := ioutil.ReadDir(filePath)
		for _, f := range files {
			AddFile(path.Join(filePath, f.Name()))
		}
		return
	case mode.IsRegular():
		break
	}

	var sid ShareID
	if existingSID, ok := preferences.FileMap[filePath]; ok {
		sid = existingSID
	} else {
		// create unique share_id
		sid = RandomShareID()
		preferences.FileMap[filePath] = sid
	}

	// read the file
	fileBytes, err := ioutil.ReadFile(filePath)
	if err != nil {
		color.Red("Cannot read file: %s", err)
		return
	}

	// create the shares
	allCloudStores := preferences.AllCloudStores()
	shares := CreateShares(fileBytes, sid, len(allCloudStores))

	// iteratively upload shares with each cloud store
	for i, cs := range allCloudStores {
		cs.Upload(shares[i])
	}

	// only save pref if it's not a .chasm
	if sid != ShareID(".chasm") {
		preferences.Save()
	}

}

// DeleteFile deletes the remote share of this path by its shareId
func DeleteFile(filePath string) {
	if !IsValidPath(filePath) {
		color.Red("Path %s is in .chasmignore. No actions will be performed.", filePath)
		return
	}

	allCloudStores := preferences.AllCloudStores()

	if sid, ok := preferences.FileMap[filePath]; ok {
		// iteratively delete shares from each cloud store
		for _, cs := range allCloudStores {
			cs.Delete(ShareID(sid))
		}

		delete(preferences.FileMap, filePath)
		preferences.Save()

		color.Green("Deleted share from all cloud stores.")
		return
	}

	color.Red("Path %s is not tracked. Cannot find share id.", filePath)
}

// Restore shares to the original files
func Restore() {
	allCloudStores := preferences.AllCloudStores()
	sharePaths := make([]string, len(allCloudStores))

	// (1) first get all shares
	for i, cs := range allCloudStores {
		sp := cs.Restore()
		if sp == "" {
			color.Red("Restore failed for %v", cs)
			return
		}
		sharePaths[i] = sp
	}

	// (2) next restore .chasm file
	chasmFileBytes := restoreShareID(ShareID(chasmPrefFile), sharePaths)

	var restoredPrefs ChasmPref
	err := json.Unmarshal(chasmFileBytes, &restoredPrefs)
	if err != nil {
		color.Red("Cannot restore chasm preferecnes file from cloud services.")
		return
	}

	// (3) finally, for the remaining files, restore and save
	for filePath, sid := range restoredPrefs.FileMap {
		fileBytes := restoreShareID(sid, sharePaths)
		os.MkdirAll(path.Dir(filePath), 0770)
		err := ioutil.WriteFile(filePath, fileBytes, 0770)
		if err != nil {
			color.Red("Error writing restored file %s: %s", filePath, err)
			return
		}
	}
	color.Green("Done. Restored all files!")
}

func restoreShareID(sid ShareID, sharePaths []string) []byte {
	fileShares := make([]Share, len(sharePaths))

	for i, sp := range sharePaths {
		file := path.Join(sp, string(sid))
		dataBytes, err := ioutil.ReadFile(file)
		if err != nil {
			color.Red("(Skipping share) Cannot read file %s: %s", file, err)
			continue
		}

		share := Share{SID: sid, Data: dataBytes}
		fileShares[i] = share
	}

	return CombineShares(fileShares)
}
<commit_msg>some more color<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"

	"github.com/fatih/color"
)

/// Chasm Types ///

// CloudStore represents an external cloud storage service that is compatible
// with Chasm
type CloudStore interface {
	Upload(share Share)
	Delete(sid ShareID)

	//Restore downloads shares to local restore path
	Restore() string

	Description() string

	Clean()
}

// ChasmPref represents user/application preferences
type ChasmPref struct {
	root string

	// the cloud services sharing across
	FolderStores []FolderStore `json:"folder_stores"`

	// the cloud services sharing across
	GDriveStores []GDriveStore `json:"gdrive_stores"`

	// the cloud services sharing across
	DropboxStores []DropboxStore `json:"dropbox_stores"`

	// maps files to their shareId
	FileMap map[string]ShareID `json:"files"`
}

// RegisteredServices counts all services
func (p ChasmPref) RegisteredServices() int {
	return len(p.FolderStores) + len(p.GDriveStores) + len(p.DropboxStores)
}

// NeedSetup checks if there are enough services to run
func (p ChasmPref) NeedSetup() bool {
	return p.RegisteredServices() < 2
}

// AllCloudStores combines all the cloud stores
func (p ChasmPref) AllCloudStores() []CloudStore {

	// adjust length for new store types
	cloudStores := make([]CloudStore, p.RegisteredServices())

	// all other cloud stores go here
	ind := 0
	for i, fs := range p.FolderStores {
		cloudStores[i] = CloudStore(fs)
		ind += 1
	}

	for j, gds := range p.GDriveStores {
		cloudStores[j+ind] = CloudStore(gds)
		ind += 1
	}

	for k, dbs := range p.DropboxStores {
		cloudStores[k+ind] = CloudStore(dbs)
		ind += 1
	}

	return cloudStores
}

// Save saves the chasm preferences
func (p ChasmPref) Save() {
	chasmFilePath := p.root + string(filepath.Separator) + chasmPrefFile
	chasmFileBytes, err := json.MarshalIndent(preferences, "", "    ")
	check(err)

	ioutil.WriteFile(chasmFilePath, chasmFileBytes, 0660)
}

/// Chasm Functions ///

var preferences ChasmPref

const chasmPrefFile = ".chasm"
const chasmIgnoreFile = ".chasmignore"

// CreateOrLoadChasmDir creates the root *chasm* folder on the system
// if it does not exist or finds an existing directory
// returns if true if newly created
func CreateOrLoadChasmDir(root string) {
	os.MkdirAll(root, 0777)

	chasmFilePath := path.Join(root, chasmPrefFile)
	chasmFileBytes, err := ioutil.ReadFile(chasmFilePath)
	if err != nil {
		color.Green("Creating new .chasm secure folder")
		preferences.FileMap = make(map[string]ShareID)
		preferences.FileMap[chasmFilePath] = ShareID(chasmPrefFile)
	} else {
		json.Unmarshal(chasmFileBytes, &preferences)
	}

	chasmIgnorePath := path.Join(root, chasmIgnoreFile)
	_, err = ioutil.ReadFile(chasmIgnorePath)
	if err != nil {
		preferences.FileMap[chasmIgnorePath] = ShareID(chasmIgnoreFile)
		// add *.DS_Store to ignore file by default
		errWrite := ioutil.WriteFile(chasmIgnorePath, []byte(".DS_Store\n"), 0777)
		if errWrite != nil {
			color.Red("Error: could not write to %s: %s", chasmFilePath, errWrite)
		}
	}

	preferences.root = root
	preferences.Save()
}

// IsValidPath checks if a file path is vaild, i.e. it doesn't match any patterns
// in the .chasmignore file
func IsValidPath(filePath string) bool {
	base := filepath.Base(filePath)
	chasmIgnorePath := path.Join(preferences.root, chasmIgnoreFile)
	chasmIgnore, err := os.Open(chasmIgnorePath)
	if err != nil {
		return true
	}

	scanner := bufio.NewScanner(chasmIgnore)
	for scanner.Scan() {
		pattern := scanner.Text()
		// if the file matches anything in .chasmignore, return false
		ok, err := filepath.Match(pattern, base)
		if ok {
			return false
		}
		if err != nil {
			fmt.Println(err)
		}
	}

	// check for errors
	if err := scanner.Err(); err != nil {
		fmt.Println(err)
	}

	return true
}

// AddFile secret shares the file, and uploads each share to corresponding services
// if the file exists already, we delete the remote share first by its shareId
func AddFile(filePath string) {
	if !IsValidPath(filePath) {
		color.Red("Path %s is in .chasmignore. No actions will be performed.", filePath)
		return
	}
	file, _ := os.Open(filePath)
	fi, err := file.Stat()
	if err != nil {
		color.Red("Cannot get file info: %s", err)
		return
	}
	switch mode := fi.Mode(); {
	case mode.IsDir():
		files, _ := ioutil.ReadDir(filePath)
		for _, f := range files {
			AddFile(path.Join(filePath, f.Name()))
		}
		return
	case mode.IsRegular():
		break
	}

	var sid ShareID
	if existingSID, ok := preferences.FileMap[filePath]; ok {
		sid = existingSID
	} else {
		// create unique share_id
		sid = RandomShareID()
		preferences.FileMap[filePath] = sid
	}

	// read the file
	fileBytes, err := ioutil.ReadFile(filePath)
	if err != nil {
		color.Red("Cannot read file: %s", err)
		return
	}

	// create the shares
	allCloudStores := preferences.AllCloudStores()
	shares := CreateShares(fileBytes, sid, len(allCloudStores))

	// iteratively upload shares with each cloud store
	for i, cs := range allCloudStores {
		cs.Upload(shares[i])
	}

	// only save pref if it's not a .chasm
	if sid != ShareID(".chasm") {
		preferences.Save()
	}

}

// DeleteFile deletes the remote share of this path by its shareId
func DeleteFile(filePath string) {
	if !IsValidPath(filePath) {
		color.Red("Path %s is in .chasmignore. No actions will be performed.", filePath)
		return
	}

	allCloudStores := preferences.AllCloudStores()

	if sid, ok := preferences.FileMap[filePath]; ok {
		// iteratively delete shares from each cloud store
		for _, cs := range allCloudStores {
			cs.Delete(ShareID(sid))
		}

		delete(preferences.FileMap, filePath)
		preferences.Save()

		color.Yellow("Deleted share from all cloud stores.")
		return
	}

	color.Red("Path %s is not tracked. Cannot find share id.", filePath)
}

// Restore shares to the original files
func Restore() {
	allCloudStores := preferences.AllCloudStores()
	sharePaths := make([]string, len(allCloudStores))

	// (1) first get all shares
	for i, cs := range allCloudStores {
		sp := cs.Restore()
		if sp == "" {
			color.Red("Restore failed for %v", cs)
			return
		}
		sharePaths[i] = sp
	}

	// (2) next restore .chasm file
	chasmFileBytes := restoreShareID(ShareID(chasmPrefFile), sharePaths)

	var restoredPrefs ChasmPref
	err := json.Unmarshal(chasmFileBytes, &restoredPrefs)
	if err != nil {
		color.Red("Cannot restore chasm preferecnes file from cloud services.")
		return
	}

	// (3) finally, for the remaining files, restore and save
	for filePath, sid := range restoredPrefs.FileMap {
		fileBytes := restoreShareID(sid, sharePaths)
		os.MkdirAll(path.Dir(filePath), 0770)
		err := ioutil.WriteFile(filePath, fileBytes, 0770)
		if err != nil {
			color.Red("Error writing restored file %s: %s", filePath, err)
			return
		}
	}
	color.Green("Done. Restored all files!")
}

func restoreShareID(sid ShareID, sharePaths []string) []byte {
	fileShares := make([]Share, len(sharePaths))

	for i, sp := range sharePaths {
		file := path.Join(sp, string(sid))
		dataBytes, err := ioutil.ReadFile(file)
		if err != nil {
			color.Red("(Skipping share) Cannot read file %s: %s", file, err)
			continue
		}

		share := Share{SID: sid, Data: dataBytes}
		fileShares[i] = share
	}

	return CombineShares(fileShares)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"io/ioutil"
	"log"
	"net"
	"os"

	"github.com/bgmerrell/gochatd/chat"
	"github.com/bgmerrell/gochatd/handlers/raw"
)

var confPath string

func init() {
	flag.StringVar(&confPath, "conf-path", "gochatd-conf.json", "Configuration file path")
}

type config struct {
	LogPath    string `json:"log_path"`
	Addr       string `json:"address"`
	MaxNameLen int    `json:"max_name_length"`
	MsgBufSize int    `json:"msg_buffer_size"`
}

func main() {
	flag.Parse()
	cfgRaw, err := ioutil.ReadFile(confPath)
	if err != nil {
		log.Fatalf("Failed to read log file (%s): %s", confPath, err)
	}
	// Just use a JSON config file to avoid 3rd party dependencies (for
	// something like ini or toml)
	cfg := config{}
	err = json.Unmarshal(cfgRaw, &cfg)
	if err != nil {
		log.Fatalf("Failed to parse log file (%s): %s", confPath, err)
	}
	chatLogFile, err := os.OpenFile(cfg.LogPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
	if err != nil {
		log.Fatalf("Failed to open chat log: %s", err)
	}
	defer chatLogFile.Close()
	cm := chat.NewChatManager(chatLogFile)
	rh := raw.NewRawHandler(cfg.MsgBufSize, cfg.MaxNameLen)
	ln, err := net.Listen("tcp", cfg.Addr)
	if err != nil {
		log.Fatal(err)
	}
	for {
		conn, err := ln.Accept()
		if err != nil {

		}
		go rh.Handle(cm, conn)
	}
}
<commit_msg>Fix race condition identified by race detector<commit_after>package main

import (
	"encoding/json"
	"flag"
	"io/ioutil"
	"log"
	"net"
	"os"

	"github.com/bgmerrell/gochatd/chat"
	"github.com/bgmerrell/gochatd/handlers/raw"
)

var confPath string

func init() {
	flag.StringVar(&confPath, "conf-path", "gochatd-conf.json", "Configuration file path")
}

type config struct {
	LogPath    string `json:"log_path"`
	Addr       string `json:"address"`
	MaxNameLen int    `json:"max_name_length"`
	MsgBufSize int    `json:"msg_buffer_size"`
}

func main() {
	flag.Parse()
	cfgRaw, err := ioutil.ReadFile(confPath)
	if err != nil {
		log.Fatalf("Failed to read log file (%s): %s", confPath, err)
	}
	// Just use a JSON config file to avoid 3rd party dependencies (for
	// something like ini or toml)
	cfg := config{}
	err = json.Unmarshal(cfgRaw, &cfg)
	if err != nil {
		log.Fatalf("Failed to parse log file (%s): %s", confPath, err)
	}
	chatLogFile, err := os.OpenFile(cfg.LogPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
	if err != nil {
		log.Fatalf("Failed to open chat log: %s", err)
	}
	defer chatLogFile.Close()
	cm := chat.NewChatManager(chatLogFile)
	ln, err := net.Listen("tcp", cfg.Addr)
	if err != nil {
		log.Fatal(err)
	}
	for {
		rh := raw.NewRawHandler(cfg.MsgBufSize, cfg.MaxNameLen)
		conn, err := ln.Accept()
		if err != nil {

		}
		go rh.Handle(cm, conn)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"path"
	"strings"

	"github.com/vito/gocart/dependency"
	"github.com/vito/gocart/dependency_fetcher"
	"github.com/vito/gocart/locker"
)

type Reason interface {
	Description() string
}

type VersionMismatch struct {
	Expected string
	Current  string
}

func (self VersionMismatch) Description() string {
	return fmt.Sprintf("version mismatch:\n  want %s\n  have %s\n", red(self.Expected), green(self.Current))
}

type DirtyState struct {
	Output string
}

func (self DirtyState) Description() string {
	output := ""

	for _, line := range strings.Split(strings.Trim(self.Output, " \n"), "\n") {
		output = output + "  " + line + "\n"
	}

	return fmt.Sprintf("dirty state:\n%s", output)
}

func check(root string) {
	dirtyDependencies := findDirtyDependencies(root)

	if len(dirtyDependencies) > 0 {
		count := 0

		for path, reason := range dirtyDependencies {
			fmt.Println(bold(path))
			fmt.Println(reason.Description())

			count++
		}

		os.Exit(1)
	}
}

func findDirtyDependencies(root string) map[string]Reason {
	requestedDependencies := loadFile(path.Join(root, CartridgeFile))
	lockedDependencies := loadFile(path.Join(root, CartridgeLockFile))

	dependencies := locker.GenerateLock(requestedDependencies, lockedDependencies)

	dirtyDependencies := make(map[string]Reason)

	for _, dep := range dependencies {
		reason := checkForDirtyState(dep)

		if reason != nil {
			dirtyDependencies[dep.FullPath(GOPATH)] = reason
		}

		for dep, reason := range findDirtyDependencies(dep.FullPath(GOPATH)) {
			dirtyDependencies[dep] = reason
		}
	}

	return dirtyDependencies
}

func checkForDirtyState(dep dependency.Dependency) Reason {
	repoPath := dep.FullPath(GOPATH)

	if _, err := os.Stat(repoPath); os.IsNotExist(err) {
		return nil
	}

	currentVersion := findCurrentVersion(dep)

	if currentVersion != dep.Version {
		return VersionMismatch{
			Expected: dep.Version,
			Current:  currentVersion,
		}
	}

	repo, err := dependency_fetcher.NewRepository(repoPath)
	if err != nil {
		fatal(err.Error())
	}

	status := repo.StatusCommand()
	status.Dir = repoPath

	output, err := status.Output()
	if err != nil {
		fatal(err.Error())
	}

	// Bazaar is bizarre
	if string(output) == "working tree is out of date, run 'bzr update'\n" {
		return nil
	}

	if len(output) != 0 {
		return DirtyState{
			Output: string(output),
		}
	}

	return nil
}
<commit_msg>Favor displaying dirty repos over mismatched shas<commit_after>package main

import (
	"fmt"
	"os"
	"path"
	"strings"

	"github.com/vito/gocart/dependency"
	"github.com/vito/gocart/dependency_fetcher"
	"github.com/vito/gocart/locker"
)

type Reason interface {
	Description() string
}

type VersionMismatch struct {
	Expected string
	Current  string
}

func (self VersionMismatch) Description() string {
	return fmt.Sprintf("version mismatch:\n  want %s\n  have %s\n", red(self.Expected), green(self.Current))
}

type DirtyState struct {
	Output string
}

func (self DirtyState) Description() string {
	output := ""

	for _, line := range strings.Split(strings.Trim(self.Output, " \n"), "\n") {
		output = output + "  " + line + "\n"
	}

	return fmt.Sprintf("dirty state:\n%s", output)
}

func check(root string) {
	dirtyDependencies := findDirtyDependencies(root)

	if len(dirtyDependencies) > 0 {
		count := 0

		for path, reason := range dirtyDependencies {
			fmt.Println(bold(path))
			fmt.Println(reason.Description())

			count++
		}

		os.Exit(1)
	}
}

func findDirtyDependencies(root string) map[string]Reason {
	requestedDependencies := loadFile(path.Join(root, CartridgeFile))
	lockedDependencies := loadFile(path.Join(root, CartridgeLockFile))

	dependencies := locker.GenerateLock(requestedDependencies, lockedDependencies)

	dirtyDependencies := make(map[string]Reason)

	for _, dep := range dependencies {
		reason := checkForDirtyState(dep)

		if reason != nil {
			dirtyDependencies[dep.FullPath(GOPATH)] = reason
		}

		for dep, reason := range findDirtyDependencies(dep.FullPath(GOPATH)) {
			dirtyDependencies[dep] = reason
		}
	}

	return dirtyDependencies
}

func checkForDirtyState(dep dependency.Dependency) Reason {
	repoPath := dep.FullPath(GOPATH)

	if _, err := os.Stat(repoPath); os.IsNotExist(err) {
		return nil
	}

	repo, err := dependency_fetcher.NewRepository(repoPath)
	if err != nil {
		fatal(err.Error())
	}

	status := repo.StatusCommand()
	status.Dir = repoPath

	output, err := status.Output()
	if err != nil {
		fatal(err.Error())
	}

	// Bazaar is bizarre
	if string(output) == "working tree is out of date, run 'bzr update'\n" {
		return nil
	}

	if len(output) != 0 {
		return DirtyState{
			Output: string(output),
		}
	}

	currentVersion := findCurrentVersion(dep)

	if currentVersion != dep.Version {
		return VersionMismatch{
			Expected: dep.Version,
			Current:  currentVersion,
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package cluster

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/db/query"
	"github.com/lxc/lxd/lxd/task"
	"github.com/lxc/lxd/shared"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
	"github.com/pkg/errors"
)

// APIHeartbeatMember contains specific cluster node info.
type APIHeartbeatMember struct {
	ID            int64     // ID field value in nodes table.
	Address       string    // Host and Port of node.
	RaftID        uint64    // ID field value in raft_nodes table, zero if non-raft node.
	RaftRole      int       // Node role in the raft cluster, from the raft_nodes table
	Raft          bool      // Deprecated, use non-zero RaftID instead to indicate raft node.
	LastHeartbeat time.Time // Last time we received a successful response from node.
	Online        bool      // Calculated from offline threshold and LastHeatbeat time.
	updated       bool      // Has node been updated during this heartbeat run. Not sent to nodes.
}

// APIHeartbeatVersion contains max versions for all nodes in cluster.
type APIHeartbeatVersion struct {
	Schema        int
	APIExtensions int
}

// APIHeartbeat contains data sent to nodes in heartbeat.
type APIHeartbeat struct {
	sync.Mutex // Used to control access to Members maps.
	cluster    *db.Cluster
	Members    map[int64]APIHeartbeatMember
	Version    APIHeartbeatVersion
	Time       time.Time

	// Indicates if heartbeat contains a fresh set of node states.
	// This can be used to indicate to the receiving node that the state is fresh enough to
	// trigger node refresh activies (such as forkdns).
	FullStateList bool
}

// Update updates an existing APIHeartbeat struct with the raft and all node states supplied.
// If allNodes provided is an empty set then this is considered a non-full state list.
func (hbState *APIHeartbeat) Update(fullStateList bool, raftNodes []db.RaftNode, allNodes []db.NodeInfo, offlineThreshold time.Duration) {
	var maxSchemaVersion, maxAPIExtensionsVersion int

	if hbState.Members == nil {
		hbState.Members = make(map[int64]APIHeartbeatMember)
	}

	// If we've been supplied a fresh set of node states, this is a full state list.
	hbState.FullStateList = fullStateList

	raftNodeMap := make(map[string]db.RaftNode)

	// Convert raftNodes to a map keyed on address for lookups later.
	for _, raftNode := range raftNodes {
		raftNodeMap[raftNode.Address] = raftNode
	}

	// Add nodes (overwrites any nodes with same ID in map with fresh data).
	for _, node := range allNodes {
		member := APIHeartbeatMember{
			ID:            node.ID,
			Address:       node.Address,
			LastHeartbeat: node.Heartbeat,
			Online:        !node.Heartbeat.Before(time.Now().Add(-offlineThreshold)),
		}

		if raftNode, exists := raftNodeMap[member.Address]; exists {
			member.Raft = true // Deprecated
			member.RaftID = raftNode.ID
			member.RaftRole = int(raftNode.Role)
			delete(raftNodeMap, member.Address) // Used to check any remaining later.
		}

		// Add to the members map using the node ID (not the Raft Node ID).
		hbState.Members[node.ID] = member

		// Keep a record of highest APIExtensions and Schema version seen in all nodes.
		if node.APIExtensions > maxAPIExtensionsVersion {
			maxAPIExtensionsVersion = node.APIExtensions
		}

		if node.Schema > maxSchemaVersion {
			maxSchemaVersion = node.Schema
		}
	}

	hbState.Version = APIHeartbeatVersion{
		Schema:        maxSchemaVersion,
		APIExtensions: maxAPIExtensionsVersion,
	}

	if len(raftNodeMap) > 0 && hbState.cluster != nil {
		hbState.cluster.Transaction(func(tx *db.ClusterTx) error {
			for addr, raftNode := range raftNodeMap {
				_, err := tx.GetPendingNodeByAddress(addr)
				if err != nil {
					logger.Errorf("Unaccounted raft node(s) not found in 'nodes' table for heartbeat: %+v", raftNode)
				}
			}

			return nil
		})
	}

	return
}

// Send sends heartbeat requests to the nodes supplied and updates heartbeat state.
func (hbState *APIHeartbeat) Send(ctx context.Context, networkCert *shared.CertInfo, serverCert *shared.CertInfo, localAddress string, nodes []db.NodeInfo, delay bool) {
	heartbeatsWg := sync.WaitGroup{}
	sendHeartbeat := func(nodeID int64, address string, delay bool, heartbeatData *APIHeartbeat) {
		defer heartbeatsWg.Done()

		if delay {
			// Spread in time by waiting up to 3s less than the interval.
			time.Sleep(time.Duration(rand.Intn((heartbeatInterval*1000)-3000)) * time.Millisecond)
		}
		logger.Debugf("Sending heartbeat to %s", address)

		// Update timestamp to current, used for time skew detection
		heartbeatData.Time = time.Now().UTC()

		err := HeartbeatNode(ctx, address, networkCert, serverCert, heartbeatData)
		if err == nil {
			heartbeatData.Lock()
			// Ensure only update nodes that exist in Members already.
			hbNode, existing := hbState.Members[nodeID]
			if !existing {
				return
			}

			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			heartbeatData.Members[nodeID] = hbNode
			heartbeatData.Unlock()
			logger.Debugf("Successful heartbeat for %s", address)
		} else {
			logger.Debugf("Failed heartbeat for %s: %v", address, err)
		}
	}

	for _, node := range nodes {
		// Special case for the local node - just record the time now.
		if node.Address == localAddress {
			hbState.Lock()
			hbNode := hbState.Members[node.ID]
			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			hbState.Members[node.ID] = hbNode
			hbState.Unlock()
			continue
		}

		// Parallelize the rest.
		heartbeatsWg.Add(1)
		go sendHeartbeat(node.ID, node.Address, delay, hbState)
	}
	heartbeatsWg.Wait()
}

// HeartbeatTask returns a task function that performs leader-initiated heartbeat
// checks against all LXD nodes in the cluster.
//
// It will update the heartbeat timestamp column of the nodes table
// accordingly, and also notify them of the current list of database nodes.
func HeartbeatTask(gateway *Gateway) (task.Func, task.Schedule) {
	// Since the database APIs are blocking we need to wrap the core logic
	// and run it in a goroutine, so we can abort as soon as the context expires.
	heartbeatWrapper := func(ctx context.Context) {
		ch := make(chan struct{})
		go func() {
			gateway.heartbeat(ctx, false)
			ch <- struct{}{}
		}()
		select {
		case <-ch:
		case <-ctx.Done():
		}
	}

	schedule := task.Every(time.Duration(heartbeatInterval) * time.Second)

	return heartbeatWrapper, schedule
}

func (g *Gateway) heartbeat(ctx context.Context, initialHeartbeat bool) {
	// Avoid concurent heartbeat loops.
	// This is possible when both the task and the out of band heartbeat
	// round from a dqlite connection both kick in at the same time.
	g.heartbeatLock.Lock()
	defer g.heartbeatLock.Unlock()

	if g.Cluster == nil || g.server == nil || g.memoryDial != nil {
		// We're not a raft node or we're not clustered
		return
	}

	raftNodes, err := g.currentRaftNodes()
	if err != nil {
		if errors.Cause(err) == ErrNotLeader {
			return
		}

		logger.Error("Failed to get current raft members", log.Ctx{"err": err})
		return
	}

	if initialHeartbeat {
		logger.Debugf("Starting heartbeat round (full update)")
	} else {
		logger.Debugf("Starting heartbeat round")
	}
	if err != nil {
		logger.Warnf("Failed to get current raft nodes: %v", err)
		return
	}

	// Replace the local raft_nodes table immediately because it
	// might miss a row containing ourselves, since we might have
	// been elected leader before the former leader had chance to
	// send us a fresh update through the heartbeat pool.
	logger.Debugf("Heartbeat updating local raft nodes to %+v", raftNodes)
	err = g.db.Transaction(func(tx *db.NodeTx) error {
		return tx.ReplaceRaftNodes(raftNodes)
	})
	if err != nil {
		logger.Warn("Failed to replace local raft nodes", log.Ctx{"err": err})
		return
	}

	var allNodes []db.NodeInfo
	var localAddress string // Address of this node
	var offlineThreshold time.Duration
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		allNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		localAddress, err = tx.GetLocalNodeAddress()
		if err != nil {
			return err
		}

		offlineThreshold, err = tx.GetNodeOfflineThreshold()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}

	// Cumulative set of node states (will be written back to database once done).
	hbState := &APIHeartbeat{cluster: g.Cluster}

	// If this leader node hasn't sent a heartbeat recently, then its node state records
	// are likely out of date, this can happen when a node becomes a leader.
	// Send stale set to all nodes in database to get a fresh set of active nodes.
	if initialHeartbeat {
		hbState.Update(false, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, false)

		// We have the latest set of node states now, lets send that state set to all nodes.
		hbState.Update(true, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, false)
	} else {
		hbState.Update(true, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, true)
	}

	// Look for any new node which appeared since sending last heartbeat.
	var currentNodes []db.NodeInfo
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		currentNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warnf("Failed to get current cluster nodes: %v", err)
		return
	}

	newNodes := []db.NodeInfo{}
	for _, currentNode := range currentNodes {
		existing := false
		for _, node := range allNodes {
			if node.Address == currentNode.Address && node.ID == currentNode.ID {
				existing = true
				break
			}
		}

		if !existing {
			// We found a new node
			allNodes = append(allNodes, currentNode)
			newNodes = append(newNodes, currentNode)
		}
	}

	// If any new nodes found, send heartbeat to just them (with full node state).
	if len(newNodes) > 0 {
		hbState.Update(true, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, newNodes, false)
	}

	// If the context has been cancelled, return immediately.
	if ctx.Err() != nil {
		logger.Debugf("Aborting heartbeat round")
		return
	}

	err = query.Retry(func() error {
		return g.Cluster.Transaction(func(tx *db.ClusterTx) error {
			hbTime := time.Now()
			for _, node := range hbState.Members {
				if !node.updated {
					continue
				}

				err := tx.SetNodeHeartbeat(node.Address, hbTime)
				if err != nil && errors.Cause(err) != db.ErrNoSuchObject {
					return errors.Wrapf(err, "Failed updating heartbeat time for member %q", node.Address)
				}
			}

			return nil
		})
	})
	if err != nil {
		logger.Error("Failed updating cluster heartbeats", log.Ctx{"err": err})
		return
	}

	// If full node state was sent and node refresh task is specified, run it async.
	if g.HeartbeatNodeHook != nil {
		go g.HeartbeatNodeHook(hbState)
	}

	// Update last leader heartbeat time so next time a full node state list can be sent (if not this time).
	logger.Debugf("Completed heartbeat round")
}

// heartbeatInterval Number of seconds to wait between to heartbeat rounds.
const heartbeatInterval = 10

// HeartbeatNode performs a single heartbeat request against the node with the given address.
func HeartbeatNode(taskCtx context.Context, address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo, heartbeatData *APIHeartbeat) error {
	logger.Debugf("Sending heartbeat request to %s", address)

	config, err := tlsClientConfig(networkCert, serverCert)
	if err != nil {
		return err
	}

	timeout := 2 * time.Second
	url := fmt.Sprintf("https://%s%s", address, databaseEndpoint)
	transport, cleanup := tlsTransport(config)
	defer cleanup()
	client := &http.Client{
		Transport: transport,
		Timeout:   timeout,
	}

	buffer := bytes.Buffer{}
	heartbeatData.Lock()
	err = json.NewEncoder(&buffer).Encode(heartbeatData)
	heartbeatData.Unlock()
	if err != nil {
		return err
	}

	request, err := http.NewRequest("PUT", url, bytes.NewReader(buffer.Bytes()))
	if err != nil {
		return err
	}
	setDqliteVersionHeader(request)

	// Use 1s later timeout to give HTTP client chance timeout with more useful info.
	ctx, cancel := context.WithTimeout(context.Background(), timeout+time.Second)
	defer cancel()
	request = request.WithContext(ctx)
	request.Close = true // Immediately close the connection after the request is done

	response, err := client.Do(request)
	if err != nil {
		return errors.Wrap(err, "failed to send HTTP request")
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP request failed: %s", response.Status)
	}

	return nil
}
<commit_msg>lxd/cluster/heartbeat: Use contextual logging<commit_after>package cluster

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/db/query"
	"github.com/lxc/lxd/lxd/task"
	"github.com/lxc/lxd/shared"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
	"github.com/pkg/errors"
)

// APIHeartbeatMember contains specific cluster node info.
type APIHeartbeatMember struct {
	ID            int64     // ID field value in nodes table.
	Address       string    // Host and Port of node.
	RaftID        uint64    // ID field value in raft_nodes table, zero if non-raft node.
	RaftRole      int       // Node role in the raft cluster, from the raft_nodes table
	Raft          bool      // Deprecated, use non-zero RaftID instead to indicate raft node.
	LastHeartbeat time.Time // Last time we received a successful response from node.
	Online        bool      // Calculated from offline threshold and LastHeatbeat time.
	updated       bool      // Has node been updated during this heartbeat run. Not sent to nodes.
}

// APIHeartbeatVersion contains max versions for all nodes in cluster.
type APIHeartbeatVersion struct {
	Schema        int
	APIExtensions int
}

// APIHeartbeat contains data sent to nodes in heartbeat.
type APIHeartbeat struct {
	sync.Mutex // Used to control access to Members maps.
	cluster    *db.Cluster
	Members    map[int64]APIHeartbeatMember
	Version    APIHeartbeatVersion
	Time       time.Time

	// Indicates if heartbeat contains a fresh set of node states.
	// This can be used to indicate to the receiving node that the state is fresh enough to
	// trigger node refresh activies (such as forkdns).
	FullStateList bool
}

// Update updates an existing APIHeartbeat struct with the raft and all node states supplied.
// If allNodes provided is an empty set then this is considered a non-full state list.
func (hbState *APIHeartbeat) Update(fullStateList bool, raftNodes []db.RaftNode, allNodes []db.NodeInfo, offlineThreshold time.Duration) {
	var maxSchemaVersion, maxAPIExtensionsVersion int

	if hbState.Members == nil {
		hbState.Members = make(map[int64]APIHeartbeatMember)
	}

	// If we've been supplied a fresh set of node states, this is a full state list.
	hbState.FullStateList = fullStateList

	raftNodeMap := make(map[string]db.RaftNode)

	// Convert raftNodes to a map keyed on address for lookups later.
	for _, raftNode := range raftNodes {
		raftNodeMap[raftNode.Address] = raftNode
	}

	// Add nodes (overwrites any nodes with same ID in map with fresh data).
	for _, node := range allNodes {
		member := APIHeartbeatMember{
			ID:            node.ID,
			Address:       node.Address,
			LastHeartbeat: node.Heartbeat,
			Online:        !node.Heartbeat.Before(time.Now().Add(-offlineThreshold)),
		}

		if raftNode, exists := raftNodeMap[member.Address]; exists {
			member.Raft = true // Deprecated
			member.RaftID = raftNode.ID
			member.RaftRole = int(raftNode.Role)
			delete(raftNodeMap, member.Address) // Used to check any remaining later.
		}

		// Add to the members map using the node ID (not the Raft Node ID).
		hbState.Members[node.ID] = member

		// Keep a record of highest APIExtensions and Schema version seen in all nodes.
		if node.APIExtensions > maxAPIExtensionsVersion {
			maxAPIExtensionsVersion = node.APIExtensions
		}

		if node.Schema > maxSchemaVersion {
			maxSchemaVersion = node.Schema
		}
	}

	hbState.Version = APIHeartbeatVersion{
		Schema:        maxSchemaVersion,
		APIExtensions: maxAPIExtensionsVersion,
	}

	if len(raftNodeMap) > 0 && hbState.cluster != nil {
		hbState.cluster.Transaction(func(tx *db.ClusterTx) error {
			for addr, raftNode := range raftNodeMap {
				_, err := tx.GetPendingNodeByAddress(addr)
				if err != nil {
					logger.Errorf("Unaccounted raft node(s) not found in 'nodes' table for heartbeat: %+v", raftNode)
				}
			}

			return nil
		})
	}

	return
}

// Send sends heartbeat requests to the nodes supplied and updates heartbeat state.
func (hbState *APIHeartbeat) Send(ctx context.Context, networkCert *shared.CertInfo, serverCert *shared.CertInfo, localAddress string, nodes []db.NodeInfo, delay bool) {
	heartbeatsWg := sync.WaitGroup{}
	sendHeartbeat := func(nodeID int64, address string, delay bool, heartbeatData *APIHeartbeat) {
		defer heartbeatsWg.Done()

		if delay {
			// Spread in time by waiting up to 3s less than the interval.
			time.Sleep(time.Duration(rand.Intn((heartbeatInterval*1000)-3000)) * time.Millisecond)
		}

		logger.Debug("Sending heartbeat", log.Ctx{"address": address})

		// Update timestamp to current, used for time skew detection
		heartbeatData.Time = time.Now().UTC()

		err := HeartbeatNode(ctx, address, networkCert, serverCert, heartbeatData)
		if err == nil {
			heartbeatData.Lock()
			// Ensure only update nodes that exist in Members already.
			hbNode, existing := hbState.Members[nodeID]
			if !existing {
				return
			}

			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			heartbeatData.Members[nodeID] = hbNode
			heartbeatData.Unlock()
			logger.Debug("Successful heartbeat", log.Ctx{"address": address})
		} else {
			logger.Debug("Failed heartbeat", log.Ctx{"address": address, "err": err})
		}
	}

	for _, node := range nodes {
		// Special case for the local node - just record the time now.
		if node.Address == localAddress {
			hbState.Lock()
			hbNode := hbState.Members[node.ID]
			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			hbState.Members[node.ID] = hbNode
			hbState.Unlock()
			continue
		}

		// Parallelize the rest.
		heartbeatsWg.Add(1)
		go sendHeartbeat(node.ID, node.Address, delay, hbState)
	}
	heartbeatsWg.Wait()
}

// HeartbeatTask returns a task function that performs leader-initiated heartbeat
// checks against all LXD nodes in the cluster.
//
// It will update the heartbeat timestamp column of the nodes table
// accordingly, and also notify them of the current list of database nodes.
func HeartbeatTask(gateway *Gateway) (task.Func, task.Schedule) {
	// Since the database APIs are blocking we need to wrap the core logic
	// and run it in a goroutine, so we can abort as soon as the context expires.
	heartbeatWrapper := func(ctx context.Context) {
		ch := make(chan struct{})
		go func() {
			gateway.heartbeat(ctx, false)
			ch <- struct{}{}
		}()
		select {
		case <-ch:
		case <-ctx.Done():
		}
	}

	schedule := task.Every(time.Duration(heartbeatInterval) * time.Second)

	return heartbeatWrapper, schedule
}

func (g *Gateway) heartbeat(ctx context.Context, initialHeartbeat bool) {
	// Avoid concurent heartbeat loops.
	// This is possible when both the task and the out of band heartbeat
	// round from a dqlite connection both kick in at the same time.
	g.heartbeatLock.Lock()
	defer g.heartbeatLock.Unlock()

	if g.Cluster == nil || g.server == nil || g.memoryDial != nil {
		// We're not a raft node or we're not clustered
		return
	}

	raftNodes, err := g.currentRaftNodes()
	if err != nil {
		if errors.Cause(err) == ErrNotLeader {
			return
		}

		logger.Error("Failed to get current raft members", log.Ctx{"err": err})
		return
	}

	if initialHeartbeat {
		logger.Debug("Starting heartbeat round (initial)")
	} else {
		logger.Debug("Starting heartbeat round")
	}

	// Replace the local raft_nodes table immediately because it
	// might miss a row containing ourselves, since we might have
	// been elected leader before the former leader had chance to
	// send us a fresh update through the heartbeat pool.
	logger.Debug("Heartbeat updating local raft members", log.Ctx{"members": raftNodes})
	err = g.db.Transaction(func(tx *db.NodeTx) error {
		return tx.ReplaceRaftNodes(raftNodes)
	})
	if err != nil {
		logger.Warn("Failed to replace local raft members", log.Ctx{"err": err})
		return
	}

	var allNodes []db.NodeInfo
	var localAddress string // Address of this node
	var offlineThreshold time.Duration
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		allNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		localAddress, err = tx.GetLocalNodeAddress()
		if err != nil {
			return err
		}

		offlineThreshold, err = tx.GetNodeOfflineThreshold()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}

	// Cumulative set of node states (will be written back to database once done).
	hbState := &APIHeartbeat{cluster: g.Cluster}

	// If this leader node hasn't sent a heartbeat recently, then its node state records
	// are likely out of date, this can happen when a node becomes a leader.
	// Send stale set to all nodes in database to get a fresh set of active nodes.
	if initialHeartbeat {
		hbState.Update(false, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, false)

		// We have the latest set of node states now, lets send that state set to all nodes.
		hbState.Update(true, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, false)
	} else {
		hbState.Update(true, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, true)
	}

	// Look for any new node which appeared since sending last heartbeat.
	var currentNodes []db.NodeInfo
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		currentNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warnf("Failed to get current cluster nodes: %v", err)
		return
	}

	newNodes := []db.NodeInfo{}
	for _, currentNode := range currentNodes {
		existing := false
		for _, node := range allNodes {
			if node.Address == currentNode.Address && node.ID == currentNode.ID {
				existing = true
				break
			}
		}

		if !existing {
			// We found a new node
			allNodes = append(allNodes, currentNode)
			newNodes = append(newNodes, currentNode)
		}
	}

	// If any new nodes found, send heartbeat to just them (with full node state).
	if len(newNodes) > 0 {
		hbState.Update(true, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, newNodes, false)
	}

	// If the context has been cancelled, return immediately.
	if ctx.Err() != nil {
		logger.Debugf("Aborting heartbeat round")
		return
	}

	err = query.Retry(func() error {
		return g.Cluster.Transaction(func(tx *db.ClusterTx) error {
			hbTime := time.Now()
			for _, node := range hbState.Members {
				if !node.updated {
					continue
				}

				err := tx.SetNodeHeartbeat(node.Address, hbTime)
				if err != nil && errors.Cause(err) != db.ErrNoSuchObject {
					return errors.Wrapf(err, "Failed updating heartbeat time for member %q", node.Address)
				}
			}

			return nil
		})
	})
	if err != nil {
		logger.Error("Failed updating cluster heartbeats", log.Ctx{"err": err})
		return
	}

	// If full node state was sent and node refresh task is specified, run it async.
	if g.HeartbeatNodeHook != nil {
		go g.HeartbeatNodeHook(hbState)
	}

	// Update last leader heartbeat time so next time a full node state list can be sent (if not this time).
	logger.Debug("Completed heartbeat round")
}

// heartbeatInterval Number of seconds to wait between to heartbeat rounds.
const heartbeatInterval = 10

// HeartbeatNode performs a single heartbeat request against the node with the given address.
func HeartbeatNode(taskCtx context.Context, address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo, heartbeatData *APIHeartbeat) error {
	logger.Debugf("Sending heartbeat request to %s", address)

	config, err := tlsClientConfig(networkCert, serverCert)
	if err != nil {
		return err
	}

	timeout := 2 * time.Second
	url := fmt.Sprintf("https://%s%s", address, databaseEndpoint)
	transport, cleanup := tlsTransport(config)
	defer cleanup()
	client := &http.Client{
		Transport: transport,
		Timeout:   timeout,
	}

	buffer := bytes.Buffer{}
	heartbeatData.Lock()
	err = json.NewEncoder(&buffer).Encode(heartbeatData)
	heartbeatData.Unlock()
	if err != nil {
		return err
	}

	request, err := http.NewRequest("PUT", url, bytes.NewReader(buffer.Bytes()))
	if err != nil {
		return err
	}
	setDqliteVersionHeader(request)

	// Use 1s later timeout to give HTTP client chance timeout with more useful info.
	ctx, cancel := context.WithTimeout(context.Background(), timeout+time.Second)
	defer cancel()
	request = request.WithContext(ctx)
	request.Close = true // Immediately close the connection after the request is done

	response, err := client.Do(request)
	if err != nil {
		return errors.Wrap(err, "failed to send HTTP request")
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP request failed: %s", response.Status)
	}

	return nil
}
<|endoftext|>