<commit_before>package hayes

//
// Pretend to be a Hayes modem.
//
// References:
// - Hayes command/error documentation:
//    http://www.messagestick.net/modem/hayes_modem.html#Introduction
// - Sounds: https://en.wikipedia.org/wiki/Precise_Tone_Plan
// - RS232: https://en.wikipedia.org/wiki/RS-232
// - Serial Programming: https://en.wikibooks.org/wiki/Serial_Programming
// - Raspberry PI lib: github.com/stianeikeland/go-rpio
//

import (
	"os"
	"os/signal"
	"fmt"
	"time"
	"net"
)

/*
#include <stdio.h>
#include <unistd.h>
#include <termios.h>
char getch(){
    char ch = 0;
    struct termios old = {0};
    fflush(stdout);
    if( tcgetattr(0, &old) < 0 ) perror("tcsetattr()");
    old.c_lflag &= ~ICANON;
    old.c_lflag &= ~ECHO;
    old.c_cc[VMIN] = 1;
    old.c_cc[VTIME] = 0;
    if( tcsetattr(0, TCSANOW, &old) < 0 ) perror("tcsetattr ICANON");
    if( read(0, &ch,1) < 0 ) perror("read()");
    old.c_lflag |= ICANON;
    old.c_lflag |= ECHO;
    if(tcsetattr(0, TCSADRAIN, &old) < 0) perror("tcsetattr ~ICANON");
    return ch;
}
*/
import "C"

////////////////////////////////////////////////////////////////////////////////////

// Command Result codes
const (
	OK            = 0
	CONNECT       = 1
	RING          = 2
	NO_CARRIER    = 3
	ERROR         = 4
	CONNECT_1200  = 5
	NO_DIALTONE   = 6
	BUSY          = 7
	NO_ANSWER     = 8
	CONNECT_2400  = 10
	CONNECT_4800  = 11
	CONNECT_9600  = 12
	CONNECT_14400 = 13
	CONNECT_19200 = 14
	CONNECT_38400 = 28
)
var status_codes = map[int]string{
	OK:            "OK",
	CONNECT:       "CONNECT",
	RING:          "RING",
	NO_CARRIER:    "NO_CARRIER",
	ERROR:         "ERROR",
	CONNECT_1200:  "CONNECT_1200",
	NO_DIALTONE:   "NO_DIALTONE",
	BUSY:          "BUSY",
	NO_ANSWER:     "NO_ANSWER",
	CONNECT_2400:  "CONNECT_2400",
	CONNECT_4800:  "CONNECT_4800",
	CONNECT_9600:  "CONNECT_9600",
	CONNECT_14400: "CONNECT_14400",
	CONNECT_19200: "CONNECT_19200",
	CONNECT_38400: "CONNECT_38400",
}

const (
	COMMANDMODE = iota
	DATAMODE
)

const OFFHOOK = false
const ONHOOK = true
const __MAX_RINGS = 5
const __CONNECT_TIMEOUT = __MAX_RINGS * 6 * time.Second

//Basic modem struct
type Modem struct {
	mode int
	onhook bool
	echo bool
	speakermode int
	volume int
	verbose bool
	quiet bool
	lastcmds []string
	lastdialed string
	r [255]byte
	curreg byte
	conn net.Conn
	pins Pins
	d [10]int
}

// Print command status, subject to quiet mode and verbose mode flags
func (m *Modem) prstatus(status int) {
	if m.quiet {
		return
	}
	if m.verbose {
		fmt.Println(status_codes[status])
	} else {
		fmt.Println(status)
	}
}

// Setup/reset modem.  Also ATZ, conveniently.
func (m *Modem) reset() (int) {
	m.onHook()
	m.lowerDSR()
	m.lowerCTS()
	m.lowerRI()

	m.echo = true		// Echo local keypresses
	m.quiet = false		// Modem offers return status
	m.verbose = true	// Text return codes
	m.volume = 1		// moderate volume
	m.speakermode = 1	// on until other modem heard
	m.lastcmds = nil
	m.lastdialed = ""
	m.curreg = 0		// current register selected (from ATSn)
	m.setupRegs()
	m.setupDebug()

	time.Sleep(500 *time.Millisecond) // Make it look good

	m.raiseDSR()
	m.raiseCTS()		// Ready for DTE to send us data
	return OK
}

// Watch a subset of pins and registers and toggle the LED as apropriate
// Must be a goroutine
func (m *Modem) handlePINs() {
	for {
		// MR LED
		if m.readCTS() && m.readDSR() {
			m.led_MR_on()
		} else {
			m.led_MR_off()
		}

		// AA LED
		if m.r[REG_AUTO_ANSWER] != 0 {
			m.led_AA_on()
		} else {
			m.led_AA_off()
		}

		// TR LED
		if m.readDTR() {
			m.led_TR_on()
		} else {
			m.led_TR_off()
		}

		// CD LED
		if m.readCD() {
			m.led_CD_on()
		} else {
			m.led_CD_off()
		}

		// OH LED
		if !m.onhook {
			m.led_OH_on()
		} else {
			m.led_OH_off()
		}

		// RI LED
		if m.readRI() {
			m.led_RD_on()
		} else {
			m.led_RD_off()
		}

		// CS LED
		if m.readCTS() {
			m.led_CS_on()
		} else {
			m.led_CS_off()
		}

		// DTR PIN
		if m.readDTR() == false && !m.onhook && m.conn != nil {
			// DTE Dropped DTR, hang up the phone if DTR is not
			// reestablished withing S25 * 1/100's of a second
			time.Sleep(time.Duration(m.r[REG_DTR_DELAY]) * 100 *
				time.Millisecond)
			if m.readDTR() == false && !m.onhook && m.conn != nil {
				m.onHook()
			}
		}

		// debug
		if m.d[1] == 2 {
			m.raiseDSR()
			m.raiseCTS()
			m.d[1] = 0
		}
		if m.d[1] == 1 {
			m.lowerDSR()
			m.lowerCTS()
			m.d[1] = 0
		}

		time.Sleep(250 * time.Millisecond)
	}
}

func (m *Modem) handleModem() {
	// Handle:
	// - passing bytes from the modem to the serial port (stdout for now)
	// - accepting incoming connections (ie, noticing the phone ringing)
	// - other housekeeping tasks (eg, clearing the ring counter)
	//
	// This must be a goroutine.

	// Clear the ring counter if there's been no rings for at least 8 seconds
	last_ring_time := time.Now()
	go func() {
		for range time.Tick(8 * time.Second) {
			if time.Since(last_ring_time) >= 8 * time.Second {
				m.r[REG_RING_COUNT] = 0
			}
		}
	}()

	l, err := net.Listen("tcp", ":20000")
	if err != nil {
		panic(err)
	}
	defer l.Close()

	for {
		conn, err := l.Accept()
		if err != nil {
			debugf("l.Accept(): %s\n", err)
			continue
		}

		if !m.onhook {	// "Busy" signal.
			conn.Close()
			continue
		}

		for i := byte(0); i < __MAX_RINGS; i++ {
			last_ring_time = time.Now()
			m.prstatus(RING)
			if !m.onhook { // computer has issued 'ATA'
				m.conn = conn
				conn = nil
				goto answered
			}

			// Simulate the "2-4" pattern for POTS ring signal (2
			// seconds of high voltage ring signal, 4 seconds
			// of silence)

			// Ring for 2s
			d := 0
			m.raiseRI()
			for m.onhook  && d < 2000 {
				time.Sleep(250 * time.Millisecond)
				d += 250
				if !m.onhook { // computer has issued 'ATA'
					m.conn = conn
					conn = nil
					goto answered
				}
			}
			m.lowerRI()

			// If Auto Answer if enabled and we've
			// exceeded the configured number of rings to
			// wait before answering, answer the call.  We
			// do this here before the 4s delay as I think
			// it feels more correct.
			if m.r[REG_AUTO_ANSWER] > 0 {
				m.r[REG_RING_COUNT]++
				if m.r[REG_RING_COUNT] >= m.r[REG_AUTO_ANSWER] {
					m.answer()
				}
			}

			// Silence for 4s
			d = 0
			for m.onhook && d < 4000 {
				time.Sleep(250 * time.Millisecond)
				d += 250
				if !m.onhook { // computer has issued 'ATA'
					m.conn = conn
					conn = nil
					goto answered
				}
			}
		}

		// At this point we've not answered and have timed out.
		if m.onhook {	// computer didn't answer
			conn.Close()
			continue
		}

	answered:
		// if we're here, the computer answered, so pass bytes
		// from the remote dialer to the serial port (for now, stdout)
		// as long as we're offhook, we're in DATA MODE and we have
		// valid carrier (m.comm != nil)
		//
		// TODO: this is line based, needed to be character based
		// TODO: Blink the RD LED somewhere in here.
		// TODO: XXX Racy -- if DTE issues ATA , causing m.onhook == true
		// before m.mode == DATAMODE, the for loop exits and hangs up.
		m.r[REG_RING_COUNT] = 0
		m.lowerRI()
		buf := make([]byte, 1)
		for !m.onhook  && m.mode == DATAMODE && m.conn != nil {
			if _, err = m.conn.Read(buf); err != nil {
				break
			}
			m.led_RD_on()
			fmt.Printf("%s", string(buf)) //  Send to DTE (stdout)
			m.led_RD_off()
		}

		// If we're here, we lost "carrier" somehow.
		m.led_RD_off()
		m.prstatus(NO_CARRIER)
		m.onHook()
		if conn != nil {
			conn.Close() // just to be safe?
		}
	}
}

// Catch ^C, reset the HW pins
func (m *Modem) signalHandler() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	// Block until a signal is received.
	s := <-c
	fmt.Println("Got signal:", s)
	m.clearPins()
	os.Exit(0)
}

// Boot the modem
func (m *Modem) PowerOn() {
	m.setupPins()		// Setup LED and GPIO pins
	m.reset()		// Setup modem inital state (or reset to
	                        // initial state)
	go m.signalHandler()

	go m.handlePINs()       // Monitor input pins & internal registers
	go m.handleModem()	// Handle in-bound in a seperate goroutine

	// Signal to DTE that we're ready
	m.raiseDSR()
	m.raiseCTS()
	// Tell user (if they're looking) we're ready
	m.prstatus(OK)

	// Consume bytes from the serial port and process or send to remote
	// as per m.mode
	var c byte
	var s string
	var lastthree [3]byte
	var out []byte
	var idx int
	var guard_time time.Duration
	var sinceLastChar time.Time

	out = make([]byte, 1)
	for {
		// XXX
		// becuse this is not just a modem  program yet, some static
		// key mapping is needed
		c = byte(C.getch())
		if c == 127 {	// ASCII DEL -> ASCII BS
			c = m.r[REG_BS_CH]
		}
		// Ignore anything above ASCII 127 or the ASCII escape
		if c > 127 || c == 27 {
			continue
		}
		// end of key mappings

		if m.echo {
			fmt.Printf("%c", c)
			// XXX: handle backspace
			if c == m.r[REG_BS_CH] {
				fmt.Printf(" %c", c)
			}
		}

		switch m.mode {
		case COMMANDMODE:
			if c == m.r[REG_CR_CH] && s != "" {
				m.command(s)
				s = ""
			}  else if c == m.r[REG_BS_CH]  && len(s) > 0 {
				s = s[0:len(s) - 1]
			} else {
				s += string(c)
			}

		case DATAMODE:
			if m.onhook == false && m.conn != nil {
				m.led_SD_on()
				out = make([]byte, 1)
				out[0] = c
				m.conn.Write(out)
				time.Sleep(10 *time.Millisecond) // HACK!
				m.led_SD_off()
				// TODO: make sure the LED says on long enough
			}

			// Look for the command escape sequence
			lastthree[idx] = c
			idx = (idx + 1) % 3
			guard_time =
				time.Duration(float64(m.r[REG_ESC_CODE_GUARD]) *
				0.02) * time.Second

			if lastthree[0] == m.r[REG_ESC_CH] &&
				lastthree[1] == m.r[REG_ESC_CH] &&
				lastthree[2] == m.r[REG_ESC_CH] &&
				time.Since(sinceLastChar) >
				time.Duration(guard_time)  {
				m.mode = COMMANDMODE
				m.prstatus(OK) // signal that we're in command mode
				continue
			}
			if c != '+' {
				sinceLastChar = time.Now()
			}
		}
	}
	m.lowerDSR()
	m.lowerCTS()
}
<commit_msg>Break out error/result codes into separate file<commit_after>package hayes

//
// Pretend to be a Hayes modem.
//
// References:
// - Hayes command/error documentation:
//    http://www.messagestick.net/modem/hayes_modem.html#Introduction
// - Sounds: https://en.wikipedia.org/wiki/Precise_Tone_Plan
// - RS232: https://en.wikipedia.org/wiki/RS-232
// - Serial Programming: https://en.wikibooks.org/wiki/Serial_Programming
// - Raspberry PI lib: github.com/stianeikeland/go-rpio
//

import (
	"os"
	"os/signal"
	"fmt"
	"time"
	"net"
)

/*
#include <stdio.h>
#include <unistd.h>
#include <termios.h>
char getch(){
    char ch = 0;
    struct termios old = {0};
    fflush(stdout);
    if( tcgetattr(0, &old) < 0 ) perror("tcsetattr()");
    old.c_lflag &= ~ICANON;
    old.c_lflag &= ~ECHO;
    old.c_cc[VMIN] = 1;
    old.c_cc[VTIME] = 0;
    if( tcsetattr(0, TCSANOW, &old) < 0 ) perror("tcsetattr ICANON");
    if( read(0, &ch,1) < 0 ) perror("read()");
    old.c_lflag |= ICANON;
    old.c_lflag |= ECHO;
    if(tcsetattr(0, TCSADRAIN, &old) < 0) perror("tcsetattr ~ICANON");
    return ch;
}
*/
import "C"

////////////////////////////////////////////////////////////////////////////////////

const (
	COMMANDMODE = iota
	DATAMODE
)

const OFFHOOK = false
const ONHOOK = true
const __MAX_RINGS = 5
const __CONNECT_TIMEOUT = __MAX_RINGS * 6 * time.Second

//Basic modem struct
type Modem struct {
	mode int
	onhook bool
	echo bool
	speakermode int
	volume int
	verbose bool
	quiet bool
	lastcmds []string
	lastdialed string
	r [255]byte
	curreg byte
	conn net.Conn
	pins Pins
	d [10]int
}

// Setup/reset modem.  Also ATZ, conveniently.
func (m *Modem) reset() (int) {
	m.onHook()
	m.lowerDSR()
	m.lowerCTS()
	m.lowerRI()

	m.echo = true		// Echo local keypresses
	m.quiet = false		// Modem offers return status
	m.verbose = true	// Text return codes
	m.volume = 1		// moderate volume
	m.speakermode = 1	// on until other modem heard
	m.lastcmds = nil
	m.lastdialed = ""
	m.curreg = 0		// current register selected (from ATSn)
	m.setupRegs()
	m.setupDebug()

	time.Sleep(500 *time.Millisecond) // Make it look good

	m.raiseDSR()
	m.raiseCTS()		// Ready for DTE to send us data
	return OK
}

// Watch a subset of pins and registers and toggle the LED as apropriate
// Must be a goroutine
func (m *Modem) handlePINs() {
	for {
		// MR LED
		if m.readCTS() && m.readDSR() {
			m.led_MR_on()
		} else {
			m.led_MR_off()
		}

		// AA LED
		if m.r[REG_AUTO_ANSWER] != 0 {
			m.led_AA_on()
		} else {
			m.led_AA_off()
		}

		// TR LED
		if m.readDTR() {
			m.led_TR_on()
		} else {
			m.led_TR_off()
		}

		// CD LED
		if m.readCD() {
			m.led_CD_on()
		} else {
			m.led_CD_off()
		}

		// OH LED
		if !m.onhook {
			m.led_OH_on()
		} else {
			m.led_OH_off()
		}

		// RI LED
		if m.readRI() {
			m.led_RD_on()
		} else {
			m.led_RD_off()
		}

		// CS LED
		if m.readCTS() {
			m.led_CS_on()
		} else {
			m.led_CS_off()
		}

		// DTR PIN
		if m.readDTR() == false && !m.onhook && m.conn != nil {
			// DTE Dropped DTR, hang up the phone if DTR is not
			// reestablished withing S25 * 1/100's of a second
			time.Sleep(time.Duration(m.r[REG_DTR_DELAY]) * 100 *
				time.Millisecond)
			if m.readDTR() == false && !m.onhook && m.conn != nil {
				m.onHook()
			}
		}

		// debug
		if m.d[1] == 2 {
			m.raiseDSR()
			m.raiseCTS()
			m.d[1] = 0
		}
		if m.d[1] == 1 {
			m.lowerDSR()
			m.lowerCTS()
			m.d[1] = 0
		}

		time.Sleep(250 * time.Millisecond)
	}
}

func (m *Modem) handleModem() {
	// Handle:
	// - passing bytes from the modem to the serial port (stdout for now)
	// - accepting incoming connections (ie, noticing the phone ringing)
	// - other housekeeping tasks (eg, clearing the ring counter)
	//
	// This must be a goroutine.

	// Clear the ring counter if there's been no rings for at least 8 seconds
	last_ring_time := time.Now()
	go func() {
		for range time.Tick(8 * time.Second) {
			if time.Since(last_ring_time) >= 8 * time.Second {
				m.r[REG_RING_COUNT] = 0
			}
		}
	}()

	l, err := net.Listen("tcp", ":20000")
	if err != nil {
		panic(err)
	}
	defer l.Close()

	for {
		conn, err := l.Accept()
		if err != nil {
			debugf("l.Accept(): %s\n", err)
			continue
		}

		if !m.onhook {	// "Busy" signal.
			conn.Close()
			continue
		}

		for i := byte(0); i < __MAX_RINGS; i++ {
			last_ring_time = time.Now()
			m.prstatus(RING)
			if !m.onhook { // computer has issued 'ATA'
				m.conn = conn
				conn = nil
				goto answered
			}

			// Simulate the "2-4" pattern for POTS ring signal (2
			// seconds of high voltage ring signal, 4 seconds
			// of silence)

			// Ring for 2s
			d := 0
			m.raiseRI()
			for m.onhook  && d < 2000 {
				time.Sleep(250 * time.Millisecond)
				d += 250
				if !m.onhook { // computer has issued 'ATA'
					m.conn = conn
					conn = nil
					goto answered
				}
			}
			m.lowerRI()

			// If Auto Answer if enabled and we've
			// exceeded the configured number of rings to
			// wait before answering, answer the call.  We
			// do this here before the 4s delay as I think
			// it feels more correct.
			if m.r[REG_AUTO_ANSWER] > 0 {
				m.r[REG_RING_COUNT]++
				if m.r[REG_RING_COUNT] >= m.r[REG_AUTO_ANSWER] {
					m.answer()
				}
			}

			// Silence for 4s
			d = 0
			for m.onhook && d < 4000 {
				time.Sleep(250 * time.Millisecond)
				d += 250
				if !m.onhook { // computer has issued 'ATA'
					m.conn = conn
					conn = nil
					goto answered
				}
			}
		}

		// At this point we've not answered and have timed out.
		if m.onhook {	// computer didn't answer
			conn.Close()
			continue
		}

	answered:
		// if we're here, the computer answered, so pass bytes
		// from the remote dialer to the serial port (for now, stdout)
		// as long as we're offhook, we're in DATA MODE and we have
		// valid carrier (m.comm != nil)
		//
		// TODO: this is line based, needed to be character based
		// TODO: Blink the RD LED somewhere in here.
		// TODO: XXX Racy -- if DTE issues ATA , causing m.onhook == true
		// before m.mode == DATAMODE, the for loop exits and hangs up.
		m.r[REG_RING_COUNT] = 0
		m.lowerRI()
		buf := make([]byte, 1)
		for !m.onhook  && m.mode == DATAMODE && m.conn != nil {
			if _, err = m.conn.Read(buf); err != nil {
				break
			}
			m.led_RD_on()
			fmt.Printf("%s", string(buf)) //  Send to DTE (stdout)
			m.led_RD_off()
		}

		// If we're here, we lost "carrier" somehow.
		m.led_RD_off()
		m.prstatus(NO_CARRIER)
		m.onHook()
		if conn != nil {
			conn.Close() // just to be safe?
		}
	}
}

// Catch ^C, reset the HW pins
func (m *Modem) signalHandler() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	// Block until a signal is received.
	s := <-c
	fmt.Println("Got signal:", s)
	m.clearPins()
	os.Exit(0)
}

// Boot the modem
func (m *Modem) PowerOn() {
	m.setupPins()		// Setup LED and GPIO pins
	m.reset()		// Setup modem inital state (or reset to
	                        // initial state)
	go m.signalHandler()

	go m.handlePINs()       // Monitor input pins & internal registers
	go m.handleModem()	// Handle in-bound in a seperate goroutine

	// Signal to DTE that we're ready
	m.raiseDSR()
	m.raiseCTS()
	// Tell user (if they're looking) we're ready
	m.prstatus(OK)

	// Consume bytes from the serial port and process or send to remote
	// as per m.mode
	var c byte
	var s string
	var lastthree [3]byte
	var out []byte
	var idx int
	var guard_time time.Duration
	var sinceLastChar time.Time

	out = make([]byte, 1)
	for {
		// XXX
		// becuse this is not just a modem  program yet, some static
		// key mapping is needed
		c = byte(C.getch())
		if c == 127 {	// ASCII DEL -> ASCII BS
			c = m.r[REG_BS_CH]
		}
		// Ignore anything above ASCII 127 or the ASCII escape
		if c > 127 || c == 27 {
			continue
		}
		// end of key mappings

		if m.echo {
			fmt.Printf("%c", c)
			// XXX: handle backspace
			if c == m.r[REG_BS_CH] {
				fmt.Printf(" %c", c)
			}
		}

		switch m.mode {
		case COMMANDMODE:
			if c == m.r[REG_CR_CH] && s != "" {
				m.command(s)
				s = ""
			}  else if c == m.r[REG_BS_CH]  && len(s) > 0 {
				s = s[0:len(s) - 1]
			} else {
				s += string(c)
			}

		case DATAMODE:
			if m.onhook == false && m.conn != nil {
				m.led_SD_on()
				out = make([]byte, 1)
				out[0] = c
				m.conn.Write(out)
				time.Sleep(10 *time.Millisecond) // HACK!
				m.led_SD_off()
				// TODO: make sure the LED says on long enough
			}

			// Look for the command escape sequence
			lastthree[idx] = c
			idx = (idx + 1) % 3
			guard_time =
				time.Duration(float64(m.r[REG_ESC_CODE_GUARD]) *
				0.02) * time.Second

			if lastthree[0] == m.r[REG_ESC_CH] &&
				lastthree[1] == m.r[REG_ESC_CH] &&
				lastthree[2] == m.r[REG_ESC_CH] &&
				time.Since(sinceLastChar) >
				time.Duration(guard_time)  {
				m.mode = COMMANDMODE
				m.prstatus(OK) // signal that we're in command mode
				continue
			}
			if c != '+' {
				sinceLastChar = time.Now()
			}
		}
	}
	m.lowerDSR()
	m.lowerCTS()
}
<|endoftext|>


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

package backend

import (
	"context"
	"fmt"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes/timestamp"

	"google.golang.org/api/googleapi"

	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/luci/appengine/tq"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"

	"go.chromium.org/luci/gce/api/config/v1"
	"go.chromium.org/luci/gce/api/tasks/v1"
	"go.chromium.org/luci/gce/appengine/backend/internal/metrics"
	"go.chromium.org/luci/gce/appengine/model"
)

// countVMsQueue is the name of the count VMs task handler queue.
const countVMsQueue = "count-vms"

// countVMs counts the VMs for a given config.
func countVMs(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.CountVMs)
	switch {
	case !ok:
		return errors.Reason("unexpected payload type %T", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	}
	// Count VMs per project, server and zone.
	// VMs created from the same config eventually have the same project, server,
	// and zone but may currently exist for a previous version of the config.
	vms := &metrics.InstanceCount{}

	// Get the configured count.
	cfg := &model.Config{
		ID: task.Id,
	}
	switch err := datastore.Get(c, cfg); {
	case err == datastore.ErrNoSuchEntity:
	case err != nil:
		return errors.Annotate(err, "failed to fetch config").Err()
	default:
		amt, err := cfg.Config.Amount.GetAmount(clock.Now(c))
		if err != nil {
			return errors.Annotate(err, "failed to parse amount").Err()
		}
		vms.AddConfigured(int(amt), cfg.Config.Attributes.Project)
	}

	// Get the actual (connected, created) counts.
	var keys []*datastore.Key
	q := datastore.NewQuery(model.VMKind).Eq("config", task.Id)
	if err := datastore.GetAll(c, q, &keys); err != nil {
		return errors.Annotate(err, "failed to fetch VMs").Err()
	}
	vm := &model.VM{}
	for _, k := range keys {
		id := k.StringID()
		vm.ID = id
		switch err := datastore.Get(c, vm); {
		case err == datastore.ErrNoSuchEntity:
		case err != nil:
			return errors.Annotate(err, "failed to fetch VM").Err()
		default:
			if vm.Created > 0 {
				vms.AddCreated(1, vm.Attributes.Project, vm.Attributes.Zone)
			}
			if vm.Connected > 0 {
				vms.AddConnected(1, vm.Attributes.Project, vm.Swarming, vm.Attributes.Zone)
			}
		}
	}
	if err := vms.Update(c, task.Id); err != nil {
		return errors.Annotate(err, "failed to update count").Err()
	}
	return nil
}

// drainVM drains a given VM if necessary.
func drainVM(c context.Context, vm *model.VM) error {
	if vm.Drained {
		return nil
	}
	cfg := &model.Config{
		ID: vm.Config,
	}
	switch err := datastore.Get(c, cfg); {
	case err == datastore.ErrNoSuchEntity:
		logging.Debugf(c, "config %q does not exist", cfg.ID)
	case err != nil:
		return errors.Annotate(err, "failed to fetch config").Err()
	}
	switch amt, err := cfg.Config.Amount.GetAmount(clock.Now(c)); {
	case err != nil:
		return errors.Annotate(err, "failed to parse amount").Err()
	case amt > vm.Index:
		return nil
	default:
		logging.Debugf(c, "config %q only specifies %d VMs", cfg.ID, amt)
	}
	return datastore.RunInTransaction(c, func(c context.Context) error {
		switch err := datastore.Get(c, vm); {
		case err == datastore.ErrNoSuchEntity:
			vm.Drained = true
			return nil
		case err != nil:
			return errors.Annotate(err, "failed to fetch VM").Err()
		case vm.Drained:
			return nil
		}
		vm.Drained = true
		if err := datastore.Put(c, vm); err != nil {
			return errors.Annotate(err, "failed to store VM").Err()
		}
		return nil
	}, nil)
}

// getSuffix returns a random suffix to use when naming a GCE instance.
func getSuffix(c context.Context) string {
	const allowed = "abcdefghijklmnopqrstuvwxyz0123456789"
	suf := make([]byte, 4)
	for i := range suf {
		suf[i] = allowed[mathrand.Intn(c, len(allowed))]
	}
	return string(suf)
}

// createVMQueue is the name of the create VM task handler queue.
const createVMQueue = "create-vm"

// createVM creates a VM if it doesn't already exist.
func createVM(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.CreateVM)
	switch {
	case !ok:
		return errors.Reason("unexpected payload type %T", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	case task.GetConfig() == "":
		return errors.Reason("config is required").Err()
	}
	vm := &model.VM{
		ID:       task.Id,
		Config:   task.Config,
		Hostname: fmt.Sprintf("%s-%d-%s", task.Prefix, task.Index, getSuffix(c)),
		Index:    task.Index,
		Lifetime: task.Lifetime,
		Prefix:   task.Prefix,
		Revision: task.Revision,
		Swarming: task.Swarming,
		Timeout:  task.Timeout,
	}
	if task.Attributes != nil {
		vm.Attributes = *task.Attributes
		// TODO(crbug/942301): Auto-select zone if zone is unspecified.
		vm.Attributes.SetZone(vm.Attributes.GetZone())
	}
	// createVM is called repeatedly, so do a fast check outside the transaction.
	// In most cases, this will skip the more expensive transactional check.
	switch err := datastore.Get(c, vm); {
	case err == datastore.ErrNoSuchEntity:
	case err != nil:
		return errors.Annotate(err, "failed to fetch VM").Err()
	default:
		return nil
	}
	return datastore.RunInTransaction(c, func(c context.Context) error {
		switch err := datastore.Get(c, vm); {
		case err == datastore.ErrNoSuchEntity:
		case err != nil:
			return errors.Annotate(err, "failed to fetch VM").Err()
		default:
			return nil
		}
		if err := datastore.Put(c, vm); err != nil {
			return errors.Annotate(err, "failed to store VM").Err()
		}
		return nil
	}, nil)
}

// expandConfigQueue is the name of the expand config task handler queue.
const expandConfigQueue = "expand-config"

// expandConfig creates task queue tasks to create each VM in the given config.
func expandConfig(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.ExpandConfig)
	switch {
	case !ok:
		return errors.Reason("unexpected payload type %T", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	}
	cfg, err := getConfig(c).Get(c, &config.GetRequest{Id: task.Id})
	if err != nil {
		return errors.Annotate(err, "failed to fetch config").Err()
	}
	now := clock.Now(c)
	amt, err := cfg.Amount.GetAmount(now)
	if err != nil {
		return errors.Annotate(err, "failed to parse amount").Err()
	}
	t := make([]*tq.Task, amt)
	for i := int32(0); i < amt; i++ {
		t[i] = &tq.Task{
			Payload: &tasks.CreateVM{
				Id:         fmt.Sprintf("%s-%d", cfg.Prefix, i),
				Attributes: cfg.Attributes,
				Config:     task.Id,
				Created: &timestamp.Timestamp{
					Seconds: now.Unix(),
				},
				Index:    i,
				Lifetime: cfg.Lifetime.GetSeconds(),
				Prefix:   cfg.Prefix,
				Revision: cfg.Revision,
				Swarming: cfg.Swarming,
				Timeout:  cfg.Timeout.GetSeconds(),
			},
		}
	}
	logging.Debugf(c, "creating %d VMs", len(t))
	if err := getDispatcher(c).AddTask(c, t...); err != nil {
		return errors.Annotate(err, "failed to schedule tasks").Err()
	}
	return nil
}

// reportQuotaQueue is the name of the report quota task handler queue.
const reportQuotaQueue = "report-quota"

// reportQuota reports GCE quota utilization.
func reportQuota(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.ReportQuota)
	switch {
	case !ok:
		return errors.Reason("unexpected payload type %T", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	}
	p := &model.Project{
		ID: task.Id,
	}
	if err := datastore.Get(c, p); err != nil {
		return errors.Annotate(err, "failed to fetch project").Err()
	}
	mets := stringset.NewFromSlice(p.Config.Metric...)
	regs := stringset.NewFromSlice(p.Config.Region...)
	rsp, err := getCompute(c).Regions.List(p.Config.Project).Context(c).Do()
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			logErrors(c, gerr)
			return errors.Reason("failed to fetch quota").Err()
		}
		return errors.Annotate(err, "failed to fetch quota").Err()
	}
	for _, r := range rsp.Items {
		if regs.Has(r.Name) {
			for _, q := range r.Quotas {
				if mets.Has(q.Metric) {
					metrics.UpdateQuota(c, q.Limit, q.Usage, q.Metric, p.Config.Project, r.Name)
				}
			}
		}
	}
	return nil
}
<commit_msg>[GCE] Use datastore.Run to count VMs to limit memory usage<commit_after>// Copyright 2018 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package backend

import (
	"context"
	"fmt"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes/timestamp"

	"google.golang.org/api/googleapi"

	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/luci/appengine/tq"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"

	"go.chromium.org/luci/gce/api/config/v1"
	"go.chromium.org/luci/gce/api/tasks/v1"
	"go.chromium.org/luci/gce/appengine/backend/internal/metrics"
	"go.chromium.org/luci/gce/appengine/model"
)

// countVMsQueue is the name of the count VMs task handler queue.
const countVMsQueue = "count-vms"

// countVMs counts the VMs for a given config.
func countVMs(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.CountVMs)
	switch {
	case !ok:
		return errors.Reason("unexpected payload type %T", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	}
	// Count VMs per project, server and zone.
	// VMs created from the same config eventually have the same project, server,
	// and zone but may currently exist for a previous version of the config.
	vms := &metrics.InstanceCount{}

	// Get the configured count.
	cfg := &model.Config{
		ID: task.Id,
	}
	switch err := datastore.Get(c, cfg); {
	case err == datastore.ErrNoSuchEntity:
	case err != nil:
		return errors.Annotate(err, "failed to fetch config").Err()
	default:
		amt, err := cfg.Config.Amount.GetAmount(clock.Now(c))
		if err != nil {
			return errors.Annotate(err, "failed to parse amount").Err()
		}
		vms.AddConfigured(int(amt), cfg.Config.Attributes.Project)
	}

	// Get the actual (connected, created) counts.
	vm := &model.VM{}
	q := datastore.NewQuery(model.VMKind).Eq("config", task.Id)
	if err := datastore.Run(c, q, func(k *datastore.Key) error {
		id := k.StringID()
		vm.ID = id
		switch err := datastore.Get(c, vm); {
		case err == datastore.ErrNoSuchEntity:
			return nil
		case err != nil:
			return errors.Annotate(err, "failed to fetch VM").Err()
		default:
			if vm.Created > 0 {
				vms.AddCreated(1, vm.Attributes.Project, vm.Attributes.Zone)
			}
			if vm.Connected > 0 {
				vms.AddConnected(1, vm.Attributes.Project, vm.Swarming, vm.Attributes.Zone)
			}
			return nil
		}
	}); err != nil {
		return errors.Annotate(err, "failed to fetch VMs").Err()
	}
	if err := vms.Update(c, task.Id); err != nil {
		return errors.Annotate(err, "failed to update count").Err()
	}
	return nil
}

// drainVM drains a given VM if necessary.
func drainVM(c context.Context, vm *model.VM) error {
	if vm.Drained {
		return nil
	}
	cfg := &model.Config{
		ID: vm.Config,
	}
	switch err := datastore.Get(c, cfg); {
	case err == datastore.ErrNoSuchEntity:
		logging.Debugf(c, "config %q does not exist", cfg.ID)
	case err != nil:
		return errors.Annotate(err, "failed to fetch config").Err()
	}
	switch amt, err := cfg.Config.Amount.GetAmount(clock.Now(c)); {
	case err != nil:
		return errors.Annotate(err, "failed to parse amount").Err()
	case amt > vm.Index:
		return nil
	default:
		logging.Debugf(c, "config %q only specifies %d VMs", cfg.ID, amt)
	}
	return datastore.RunInTransaction(c, func(c context.Context) error {
		switch err := datastore.Get(c, vm); {
		case err == datastore.ErrNoSuchEntity:
			vm.Drained = true
			return nil
		case err != nil:
			return errors.Annotate(err, "failed to fetch VM").Err()
		case vm.Drained:
			return nil
		}
		vm.Drained = true
		if err := datastore.Put(c, vm); err != nil {
			return errors.Annotate(err, "failed to store VM").Err()
		}
		return nil
	}, nil)
}

// getSuffix returns a random suffix to use when naming a GCE instance.
func getSuffix(c context.Context) string {
	const allowed = "abcdefghijklmnopqrstuvwxyz0123456789"
	suf := make([]byte, 4)
	for i := range suf {
		suf[i] = allowed[mathrand.Intn(c, len(allowed))]
	}
	return string(suf)
}

// createVMQueue is the name of the create VM task handler queue.
const createVMQueue = "create-vm"

// createVM creates a VM if it doesn't already exist.
func createVM(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.CreateVM)
	switch {
	case !ok:
		return errors.Reason("unexpected payload type %T", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	case task.GetConfig() == "":
		return errors.Reason("config is required").Err()
	}
	vm := &model.VM{
		ID:       task.Id,
		Config:   task.Config,
		Hostname: fmt.Sprintf("%s-%d-%s", task.Prefix, task.Index, getSuffix(c)),
		Index:    task.Index,
		Lifetime: task.Lifetime,
		Prefix:   task.Prefix,
		Revision: task.Revision,
		Swarming: task.Swarming,
		Timeout:  task.Timeout,
	}
	if task.Attributes != nil {
		vm.Attributes = *task.Attributes
		// TODO(crbug/942301): Auto-select zone if zone is unspecified.
		vm.Attributes.SetZone(vm.Attributes.GetZone())
	}
	// createVM is called repeatedly, so do a fast check outside the transaction.
	// In most cases, this will skip the more expensive transactional check.
	switch err := datastore.Get(c, vm); {
	case err == datastore.ErrNoSuchEntity:
	case err != nil:
		return errors.Annotate(err, "failed to fetch VM").Err()
	default:
		return nil
	}
	return datastore.RunInTransaction(c, func(c context.Context) error {
		switch err := datastore.Get(c, vm); {
		case err == datastore.ErrNoSuchEntity:
		case err != nil:
			return errors.Annotate(err, "failed to fetch VM").Err()
		default:
			return nil
		}
		if err := datastore.Put(c, vm); err != nil {
			return errors.Annotate(err, "failed to store VM").Err()
		}
		return nil
	}, nil)
}

// expandConfigQueue is the name of the expand config task handler queue.
const expandConfigQueue = "expand-config"

// expandConfig creates task queue tasks to create each VM in the given config.
func expandConfig(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.ExpandConfig)
	switch {
	case !ok:
		return errors.Reason("unexpected payload type %T", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	}
	cfg, err := getConfig(c).Get(c, &config.GetRequest{Id: task.Id})
	if err != nil {
		return errors.Annotate(err, "failed to fetch config").Err()
	}
	now := clock.Now(c)
	amt, err := cfg.Amount.GetAmount(now)
	if err != nil {
		return errors.Annotate(err, "failed to parse amount").Err()
	}
	t := make([]*tq.Task, amt)
	for i := int32(0); i < amt; i++ {
		t[i] = &tq.Task{
			Payload: &tasks.CreateVM{
				Id:         fmt.Sprintf("%s-%d", cfg.Prefix, i),
				Attributes: cfg.Attributes,
				Config:     task.Id,
				Created: &timestamp.Timestamp{
					Seconds: now.Unix(),
				},
				Index:    i,
				Lifetime: cfg.Lifetime.GetSeconds(),
				Prefix:   cfg.Prefix,
				Revision: cfg.Revision,
				Swarming: cfg.Swarming,
				Timeout:  cfg.Timeout.GetSeconds(),
			},
		}
	}
	logging.Debugf(c, "creating %d VMs", len(t))
	if err := getDispatcher(c).AddTask(c, t...); err != nil {
		return errors.Annotate(err, "failed to schedule tasks").Err()
	}
	return nil
}

// reportQuotaQueue is the name of the report quota task handler queue.
const reportQuotaQueue = "report-quota"

// reportQuota reports GCE quota utilization.
func reportQuota(c context.Context, payload proto.Message) error {
	task, ok := payload.(*tasks.ReportQuota)
	switch {
	case !ok:
		return errors.Reason("unexpected payload type %T", payload).Err()
	case task.GetId() == "":
		return errors.Reason("ID is required").Err()
	}
	p := &model.Project{
		ID: task.Id,
	}
	if err := datastore.Get(c, p); err != nil {
		return errors.Annotate(err, "failed to fetch project").Err()
	}
	mets := stringset.NewFromSlice(p.Config.Metric...)
	regs := stringset.NewFromSlice(p.Config.Region...)
	rsp, err := getCompute(c).Regions.List(p.Config.Project).Context(c).Do()
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			logErrors(c, gerr)
			return errors.Reason("failed to fetch quota").Err()
		}
		return errors.Annotate(err, "failed to fetch quota").Err()
	}
	for _, r := range rsp.Items {
		if regs.Has(r.Name) {
			for _, q := range r.Quotas {
				if mets.Has(q.Metric) {
					metrics.UpdateQuota(c, q.Limit, q.Usage, q.Metric, p.Config.Project, r.Name)
				}
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"html/template"
	"net/http"
	"os"

	"github.com/gorilla/context"
	"github.com/gorilla/pat"
	"github.com/markbates/goth"
	"github.com/markbates/goth/gothic"
	"github.com/markbates/goth/providers/facebook"
	"github.com/markbates/goth/providers/twitter"
)

func main() {
	goth.UseProviders(
		twitter.New(os.Getenv("TWITTER_KEY"), os.Getenv("TWITTER_SECRET"), "http://localhost:3000/auth/twitter/callback"),
		facebook.New(os.Getenv("FACEBOOK_KEY"), os.Getenv("FACEBOOK_SECRET"), "http://localhost:3000/auth/facebook/callback"),
	)

	p := pat.New()
	p.Get("/auth/{provider}/callback", func(res http.ResponseWriter, req *http.Request) {
		user, err := gothic.CompleteUserAuth(res, req)
		if err != nil {
			fmt.Fprintln(res, err)
			return
		}
		t, _ := template.New("foo").Parse(userTemplate)
		t.Execute(res, user)
	})

	p.Get("/auth/{provider}", gothic.BeginAuthHandler)
	p.Get("/", func(res http.ResponseWriter, req *http.Request) {
		t, _ := template.New("foo").Parse(indexTemplate)
		t.Execute(res, nil)
	})
	http.ListenAndServe(":3000", context.ClearHandler(p))
}

var indexTemplate = `

Log in with Twitter

Log in with Facebook

`

var userTemplate = `

Name: {{.Name}}

Email: {{.Email}}

NickName: {{.NickName}}

Location: {{.Location}}

AvatarURL: {{.AvatarURL}}

Description: {{.Description}}

UserID: {{.UserID}}

AccessToken: {{.AccessToken}}

`
<commit_msg>Got rid of dependency on gorilla/context for the example app<commit_after>package main

import (
	"fmt"
	"html/template"
	"net/http"
	"os"
	"github.com/gorilla/pat"
	"github.com/markbates/goth"
	"github.com/markbates/goth/gothic"
	"github.com/markbates/goth/providers/facebook"
	"github.com/markbates/goth/providers/twitter"
)

func main() {
	goth.UseProviders(
		twitter.New(os.Getenv("TWITTER_KEY"), os.Getenv("TWITTER_SECRET"), "http://localhost:3000/auth/twitter/callback"),
		facebook.New(os.Getenv("FACEBOOK_KEY"), os.Getenv("FACEBOOK_SECRET"), "http://localhost:3000/auth/facebook/callback"),
	)

	p := pat.New()
	p.Get("/auth/{provider}/callback", func(res http.ResponseWriter, req *http.Request) {
		user, err := gothic.CompleteUserAuth(res, req)
		if err != nil {
			fmt.Fprintln(res, err)
			return
		}
		t, _ := template.New("foo").Parse(userTemplate)
		t.Execute(res, user)
	})

	p.Get("/auth/{provider}", gothic.BeginAuthHandler)
	p.Get("/", func(res http.ResponseWriter, req *http.Request) {
		t, _ := template.New("foo").Parse(indexTemplate)
		t.Execute(res, nil)
	})
	http.ListenAndServe(":3000", p)
}

var indexTemplate = `

Log in with Twitter

Log in with Facebook

`

var userTemplate = `

Name: {{.Name}}

Email: {{.Email}}

NickName: {{.NickName}}

Location: {{.Location}}

AvatarURL: {{.AvatarURL}}

Description: {{.Description}}

UserID: {{.UserID}}

AccessToken: {{.AccessToken}}

`
<|endoftext|>


<commit_before>package complaints

/* Common code for pulling out a user session cookie, populating a Context, etc.
 * Users that aren't logged in will be redirected to the specified URL.

func init() {
  http.HandleFunc("/deb", HandleWithSession(debHandler, "/")) // If no cookie, redirects to "/"
}

func debHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	sesh,ok := GetUserSession(ctx)
	str := fmt.Sprintf("OK\nemail=%s, ok=%v\n", sesh.Email, ok)
	w.Header().Set("Content-Type", "text/plain")
	w.Write([]byte(str))
}

 */

import(
	"time"
	"net/http"
	"net/http/httputil"

	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/log"

	"github.com/skypies/complaints/sessions"
)

// Pretty much all handlers should expect to be able to pluck this object out of their Context
type UserSession struct {
	Email        string     // case sensitive, sadly
	CreatedAt    time.Time  // when the user last went through the OAuth2 dance
	hasCreatedAt bool
}

func (us UserSession)HasCreatedAt() bool { return us.hasCreatedAt }

// To prevent other libs colliding in the context.Value keyspace, use this private key
type contextKey int
const sessionEmailKey contextKey = 0

type baseHandler    func(http.ResponseWriter, *http.Request)
type contextHandler func(context.Context, http.ResponseWriter, *http.Request)

func HandleWithSession(ch contextHandler, ifNoSessionRedirectTo string) baseHandler {
	return func(w http.ResponseWriter, r *http.Request) {
		ctx,_ := context.WithTimeout(appengine.NewContext(r), 55 * time.Second)

		session := sessions.Get(r)
		if session.Values["email"] == nil {
			reqBytes,_ := httputil.DumpRequest(r, true)
			log.Errorf(ctx, "session was empty; no cookie ?")
			log.Errorf(ctx, "session: %#v", session)
			for _,c := range r.Cookies() {
				log.Errorf(ctx, "cookie: %s", c)
			}
			log.Errorf(ctx, "req: %s", reqBytes)

			// If we have a URL to redirect to, in cases of no session, then do it
			if ifNoSessionRedirectTo != "" {
				http.Redirect(w, r, ifNoSessionRedirectTo, http.StatusFound)
				return
			}

		} else {
			sesh := UserSession{Email: session.Values["email"].(string)}

			if session.Values["tstamp"] != nil {
				tstampStr := session.Values["tstamp"].(string)
				tstamp,_ := time.Parse(time.RFC3339, tstampStr)
				sesh.CreatedAt = tstamp
				sesh.hasCreatedAt = true // time.IsZero seems useless
			}

			ctx = context.WithValue(ctx, sessionEmailKey, sesh)
		}

		// Call the underlying handler, with our shiny context
		ch(ctx, w, r)
	}
}

// Underlying handlers should call this to get their session object
func GetUserSession(ctx context.Context) (UserSession,bool) {
	sesh, ok := ctx.Value(sessionEmailKey).(UserSession)
	return sesh, ok
}
<commit_msg>Try and get robot uptime checks out of the error logs<commit_after>package complaints

/* Common code for pulling out a user session cookie, populating a Context, etc.
 * Users that aren't logged in will be redirected to the specified URL.

func init() {
  http.HandleFunc("/deb", HandleWithSession(debHandler, "/")) // If no cookie, redirects to "/"
}

func debHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	sesh,ok := GetUserSession(ctx)
	str := fmt.Sprintf("OK\nemail=%s, ok=%v\n", sesh.Email, ok)
	w.Header().Set("Content-Type", "text/plain")
	w.Write([]byte(str))
}

 */

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

	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/log"

	"github.com/skypies/complaints/sessions"
)

// Pretty much all handlers should expect to be able to pluck this object out of their Context
type UserSession struct {
	Email        string     // case sensitive, sadly
	CreatedAt    time.Time  // when the user last went through the OAuth2 dance
	hasCreatedAt bool
}

func (us UserSession)HasCreatedAt() bool { return us.hasCreatedAt }

// To prevent other libs colliding in the context.Value keyspace, use this private key
type contextKey int
const sessionEmailKey contextKey = 0

type baseHandler    func(http.ResponseWriter, *http.Request)
type contextHandler func(context.Context, http.ResponseWriter, *http.Request)

func HandleWithSession(ch contextHandler, ifNoSessionRedirectTo string) baseHandler {
	return func(w http.ResponseWriter, r *http.Request) {
		ctx,_ := context.WithTimeout(appengine.NewContext(r), 55 * time.Second)

		session := sessions.Get(r)

		if strings.HasPrefix(r.UserAgent(), "Google") {
			// Robot - do nothing

		} else if session.Values["email"] == nil {
			reqBytes,_ := httputil.DumpRequest(r, true)
			log.Errorf(ctx, "session was empty; no cookie ?")
			log.Errorf(ctx, "session: %#v", session)
			for _,c := range r.Cookies() {
				log.Errorf(ctx, "cookie: %s", c)
			}
			log.Errorf(ctx, "req: %s", reqBytes)

			// If we have a URL to redirect to, in cases of no session, then do it
			if ifNoSessionRedirectTo != "" {
				http.Redirect(w, r, ifNoSessionRedirectTo, http.StatusFound)
				return
			}

		} else {
			sesh := UserSession{Email: session.Values["email"].(string)}

			if session.Values["tstamp"] != nil {
				tstampStr := session.Values["tstamp"].(string)
				tstamp,_ := time.Parse(time.RFC3339, tstampStr)
				sesh.CreatedAt = tstamp
				sesh.hasCreatedAt = true // time.IsZero seems useless
			}

			ctx = context.WithValue(ctx, sessionEmailKey, sesh)
		}

		// Call the underlying handler, with our shiny context
		ch(ctx, w, r)
	}
}

// Underlying handlers should call this to get their session object
func GetUserSession(ctx context.Context) (UserSession,bool) {
	sesh, ok := ctx.Value(sessionEmailKey).(UserSession)
	return sesh, ok
}
<|endoftext|>


<commit_before>package wmi

import (
	"errors"
	"fmt"
	"log"
	"os"
	"reflect"

	ole "github.com/mattn/go-ole"
	"github.com/mattn/go-ole/oleutil"
)

var l = log.New(os.Stdout, "", log.LstdFlags)

func init() {
	ole.CoInitializeEx(0, 0)
	// todo: determine when/if to call ole.CoUninitialize()
}

var (
	ErrInvalidEntityType = errors.New("wmi: invalid entity type")
)

func Query(query string, dst interface{}) error {
	dv := reflect.ValueOf(dst)
	if dv.Kind() != reflect.Ptr || dv.IsNil() {
		return ErrInvalidEntityType
	}
	dv = dv.Elem()
	mat, elemType := checkMultiArg(dv)
	if mat == multiArgTypeInvalid {
		return ErrInvalidEntityType
	}

	unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator")
	if err != nil {
		return err
	}
	defer unknown.Release()

	wmi, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return err
	}
	defer wmi.Release()

	// service is a SWbemServices
	serviceRaw, err := oleutil.CallMethod(wmi, "ConnectServer")
	if err != nil {
		return err
	}
	service := serviceRaw.ToIDispatch()
	defer service.Release()

	// result is a SWBemObjectSet
	resultRaw, err := oleutil.CallMethod(service, "ExecQuery", query)
	if err != nil {
		return err
	}
	result := resultRaw.ToIDispatch()
	defer result.Release()

	count, err := oleInt64(result, "Count")
	if err != nil {
		return err
	}

	var errFieldMismatch error
	for i := int64(0); i < count; i++ {
		err := func() error {
			// item is a SWbemObject, but really a Win32_Process
			itemRaw, err := oleutil.CallMethod(result, "ItemIndex", i)
			if err != nil {
				return err
			}
			item := itemRaw.ToIDispatch()
			defer item.Release()

			ev := reflect.New(elemType)
			if err = loadEntity(ev.Interface(), item); err != nil {
				if _, ok := err.(*ErrFieldMismatch); ok {
					// We continue loading entities even in the face of field mismatch errors.
					// If we encounter any other error, that other error is returned. Otherwise,
					// an ErrFieldMismatch is returned.
					errFieldMismatch = err
				} else {
					return err
				}
			}
			if mat != multiArgTypeStructPtr {
				ev = ev.Elem()
			}
			dv.Set(reflect.Append(dv, ev))
			return nil
		}()
		if err != nil {
			return err
		}
	}
	return errFieldMismatch
}

// ErrFieldMismatch is returned when a field is to be loaded into a different
// type than the one it was stored from, or when a field is missing or
// unexported in the destination struct.
// StructType is the type of the struct pointed to by the destination argument.
type ErrFieldMismatch struct {
	StructType reflect.Type
	FieldName  string
	Reason     string
}

func (e *ErrFieldMismatch) Error() string {
	return fmt.Sprintf("wmi: cannot load field %q into a %q: %s",
		e.FieldName, e.StructType, e.Reason)
}

// loadEntity loads a SWbemObject into a struct pointer.
func loadEntity(dst interface{}, src *ole.IDispatch) error {
	v := reflect.ValueOf(dst).Elem()
	for i := 0; i < v.NumField(); i++ {
		f := v.Field(i)
		n := v.Type().Field(i).Name
		if !f.CanSet() {
			return &ErrFieldMismatch{
				StructType: f.Type(),
				FieldName:  n,
				Reason:     "CanSet() is false",
			}
		}
		prop, err := oleutil.GetProperty(src, n)
		if err != nil {
			return err
		}
		switch f.Kind() {
		case reflect.String:
			f.SetString(prop.ToString())
		default:
			l.Println("ignore:", n, f.Type())
		}
	}
	return nil
}

type multiArgType int

const (
	multiArgTypeInvalid multiArgType = iota
	multiArgTypeStruct
	multiArgTypeStructPtr
)

// checkMultiArg checks that v has type []S, []*S for some struct type S.
//
// It returns what category the slice's elements are, and the reflect.Type
// that represents S.
func checkMultiArg(v reflect.Value) (m multiArgType, elemType reflect.Type) {
	if v.Kind() != reflect.Slice {
		return multiArgTypeInvalid, nil
	}
	elemType = v.Type().Elem()
	switch elemType.Kind() {
	case reflect.Struct:
		return multiArgTypeStruct, elemType
	case reflect.Ptr:
		elemType = elemType.Elem()
		if elemType.Kind() == reflect.Struct {
			return multiArgTypeStructPtr, elemType
		}
	}
	return multiArgTypeInvalid, nil
}

func oleInt64(item *ole.IDispatch, prop string) (int64, error) {
	v, err := oleutil.GetProperty(item, prop)
	if err != nil {
		return 0, err
	}
	i := int64(v.Val)
	return i, nil
}
<commit_msg>Lack of property isn't an error<commit_after>package wmi

import (
	"errors"
	"fmt"
	"log"
	"os"
	"reflect"

	ole "github.com/mattn/go-ole"
	"github.com/mattn/go-ole/oleutil"
)

var l = log.New(os.Stdout, "", log.LstdFlags)

func init() {
	ole.CoInitializeEx(0, 0)
	// todo: determine when/if to call ole.CoUninitialize()
}

var (
	ErrInvalidEntityType = errors.New("wmi: invalid entity type")
)

func Query(query string, dst interface{}) error {
	dv := reflect.ValueOf(dst)
	if dv.Kind() != reflect.Ptr || dv.IsNil() {
		return ErrInvalidEntityType
	}
	dv = dv.Elem()
	mat, elemType := checkMultiArg(dv)
	if mat == multiArgTypeInvalid {
		return ErrInvalidEntityType
	}

	unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator")
	if err != nil {
		return err
	}
	defer unknown.Release()

	wmi, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return err
	}
	defer wmi.Release()

	// service is a SWbemServices
	serviceRaw, err := oleutil.CallMethod(wmi, "ConnectServer")
	if err != nil {
		return err
	}
	service := serviceRaw.ToIDispatch()
	defer service.Release()

	// result is a SWBemObjectSet
	resultRaw, err := oleutil.CallMethod(service, "ExecQuery", query)
	if err != nil {
		return err
	}
	result := resultRaw.ToIDispatch()
	defer result.Release()

	count, err := oleInt64(result, "Count")
	if err != nil {
		return err
	}

	var errFieldMismatch error
	for i := int64(0); i < count; i++ {
		err := func() error {
			// item is a SWbemObject, but really a Win32_Process
			itemRaw, err := oleutil.CallMethod(result, "ItemIndex", i)
			if err != nil {
				return err
			}
			item := itemRaw.ToIDispatch()
			defer item.Release()

			ev := reflect.New(elemType)
			if err = loadEntity(ev.Interface(), item); err != nil {
				if _, ok := err.(*ErrFieldMismatch); ok {
					// We continue loading entities even in the face of field mismatch errors.
					// If we encounter any other error, that other error is returned. Otherwise,
					// an ErrFieldMismatch is returned.
					errFieldMismatch = err
				} else {
					return err
				}
			}
			if mat != multiArgTypeStructPtr {
				ev = ev.Elem()
			}
			dv.Set(reflect.Append(dv, ev))
			return nil
		}()
		if err != nil {
			return err
		}
	}
	return errFieldMismatch
}

// ErrFieldMismatch is returned when a field is to be loaded into a different
// type than the one it was stored from, or when a field is missing or
// unexported in the destination struct.
// StructType is the type of the struct pointed to by the destination argument.
type ErrFieldMismatch struct {
	StructType reflect.Type
	FieldName  string
	Reason     string
}

func (e *ErrFieldMismatch) Error() string {
	return fmt.Sprintf("wmi: cannot load field %q into a %q: %s",
		e.FieldName, e.StructType, e.Reason)
}

// loadEntity loads a SWbemObject into a struct pointer.
func loadEntity(dst interface{}, src *ole.IDispatch) error {
	v := reflect.ValueOf(dst).Elem()
	for i := 0; i < v.NumField(); i++ {
		f := v.Field(i)
		n := v.Type().Field(i).Name
		if !f.CanSet() {
			return &ErrFieldMismatch{
				StructType: f.Type(),
				FieldName:  n,
				Reason:     "CanSet() is false",
			}
		}
		prop, err := oleutil.GetProperty(src, n)
		if err != nil {
			continue
		}
		switch f.Kind() {
		case reflect.String:
			f.SetString(prop.ToString())
		default:
			l.Println("ignore:", n, f.Type())
		}
	}
	return nil
}

type multiArgType int

const (
	multiArgTypeInvalid multiArgType = iota
	multiArgTypeStruct
	multiArgTypeStructPtr
)

// checkMultiArg checks that v has type []S, []*S for some struct type S.
//
// It returns what category the slice's elements are, and the reflect.Type
// that represents S.
func checkMultiArg(v reflect.Value) (m multiArgType, elemType reflect.Type) {
	if v.Kind() != reflect.Slice {
		return multiArgTypeInvalid, nil
	}
	elemType = v.Type().Elem()
	switch elemType.Kind() {
	case reflect.Struct:
		return multiArgTypeStruct, elemType
	case reflect.Ptr:
		elemType = elemType.Elem()
		if elemType.Kind() == reflect.Struct {
			return multiArgTypeStructPtr, elemType
		}
	}
	return multiArgTypeInvalid, nil
}

func oleInt64(item *ole.IDispatch, prop string) (int64, error) {
	v, err := oleutil.GetProperty(item, prop)
	if err != nil {
		return 0, err
	}
	i := int64(v.Val)
	return i, nil
}
<|endoftext|>


<commit_before>package handlers

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"time"
)

type Session struct {
	SessionId     int64
	SessionSecret string `json:"-"`
	UserId        int64
}

func (s *Session) Write(w http.ResponseWriter) error {
	enc := json.NewEncoder(w)
	return enc.Encode(s)
}

func (s *Session) Read(json_str string) error {
	dec := json.NewDecoder(strings.NewReader(json_str))
	return dec.Decode(s)
}

func GetSession(tx *Tx, r *http.Request) (*Session, error) {
	var s Session

	cookie, err := r.Cookie("moneygo-session")
	if err != nil {
		return nil, fmt.Errorf("moneygo-session cookie not set")
	}
	s.SessionSecret = cookie.Value

	err = tx.SelectOne(&s, "SELECT * from sessions where SessionSecret=?", s.SessionSecret)
	if err != nil {
		return nil, err
	}
	return &s, nil
}

func DeleteSessionIfExists(tx *Tx, r *http.Request) error {
	session, err := GetSession(tx, r)
	if err == nil {
		_, err := tx.Delete(session)
		if err != nil {
			return err
		}
	}
	return nil
}

func NewSessionCookie() (string, error) {
	bits := make([]byte, 128)
	if _, err := io.ReadFull(rand.Reader, bits); err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(bits), nil
}

type NewSessionWriter struct {
	session *Session
	cookie  *http.Cookie
}

func (n *NewSessionWriter) Write(w http.ResponseWriter) error {
	http.SetCookie(w, n.cookie)
	return n.session.Write(w)
}

func NewSession(tx *Tx, r *http.Request, userid int64) (*NewSessionWriter, error) {
	s := Session{}

	session_secret, err := NewSessionCookie()
	if err != nil {
		return nil, err
	}

	existing, err := tx.SelectInt("SELECT count(*) from sessions where SessionSecret=?", session_secret)
	if err != nil {
		return nil, err
	}
	if existing > 0 {
		return nil, fmt.Errorf("%d session(s) exist with the generated session_secret")
	}

	cookie := http.Cookie{
		Name:     "moneygo-session",
		Value:    session_secret,
		Path:     "/",
		Domain:   r.URL.Host,
		Expires:  time.Now().AddDate(0, 1, 0), // a month from now
		Secure:   true,
		HttpOnly: true,
	}

	s.SessionSecret = session_secret
	s.UserId = userid

	err = tx.Insert(&s)
	if err != nil {
		return nil, err
	}
	return &NewSessionWriter{&s, &cookie}, nil
}

func SessionHandler(r *http.Request, context *Context) ResponseWriterWriter {
	if r.Method == "POST" || r.Method == "PUT" {
		var user User
		if err := ReadJSON(r, &user); err != nil {
			return NewError(3 /*Invalid Request*/)
		}

		dbuser, err := GetUserByUsername(context.Tx, user.Username)
		if err != nil {
			return NewError(2 /*Unauthorized Access*/)
		}

		user.HashPassword()
		if user.PasswordHash != dbuser.PasswordHash {
			return NewError(2 /*Unauthorized Access*/)
		}

		err = DeleteSessionIfExists(context.Tx, r)
		if err != nil {
			log.Print(err)
			return NewError(999 /*Internal Error*/)
		}

		sessionwriter, err := NewSession(context.Tx, r, dbuser.UserId)
		if err != nil {
			log.Print(err)
			return NewError(999 /*Internal Error*/)
		}
		return sessionwriter
	} else if r.Method == "GET" {
		s, err := GetSession(context.Tx, r)
		if err != nil {
			return NewError(1 /*Not Signed In*/)
		}

		return s
	} else if r.Method == "DELETE" {
		err := DeleteSessionIfExists(context.Tx, r)
		if err != nil {
			log.Print(err)
			return NewError(999 /*Internal Error*/)
		}
		return SuccessWriter{}
	}
	return NewError(3 /*Invalid Request*/)
}
<commit_msg>Add expiration and creation times to sessions<commit_after>package handlers

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"time"
)

type Session struct {
	SessionId     int64
	SessionSecret string `json:"-"`
	UserId        int64
	Created       time.Time
	Expires       time.Time
}

func (s *Session) Write(w http.ResponseWriter) error {
	enc := json.NewEncoder(w)
	return enc.Encode(s)
}

func (s *Session) Read(json_str string) error {
	dec := json.NewDecoder(strings.NewReader(json_str))
	return dec.Decode(s)
}

func GetSession(tx *Tx, r *http.Request) (*Session, error) {
	var s Session

	cookie, err := r.Cookie("moneygo-session")
	if err != nil {
		return nil, fmt.Errorf("moneygo-session cookie not set")
	}
	s.SessionSecret = cookie.Value

	err = tx.SelectOne(&s, "SELECT * from sessions where SessionSecret=?", s.SessionSecret)
	if err != nil {
		return nil, err
	}

	if s.Expires.Before(time.Now()) {
		tx.Delete(&s)
		return nil, fmt.Errorf("Session has expired")
	}
	return &s, nil
}

func DeleteSessionIfExists(tx *Tx, r *http.Request) error {
	session, err := GetSession(tx, r)
	if err == nil {
		_, err := tx.Delete(session)
		if err != nil {
			return err
		}
	}
	return nil
}

func NewSessionCookie() (string, error) {
	bits := make([]byte, 128)
	if _, err := io.ReadFull(rand.Reader, bits); err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(bits), nil
}

type NewSessionWriter struct {
	session *Session
	cookie  *http.Cookie
}

func (n *NewSessionWriter) Write(w http.ResponseWriter) error {
	http.SetCookie(w, n.cookie)
	return n.session.Write(w)
}

func NewSession(tx *Tx, r *http.Request, userid int64) (*NewSessionWriter, error) {
	s := Session{}

	session_secret, err := NewSessionCookie()
	if err != nil {
		return nil, err
	}

	existing, err := tx.SelectInt("SELECT count(*) from sessions where SessionSecret=?", session_secret)
	if err != nil {
		return nil, err
	}
	if existing > 0 {
		return nil, fmt.Errorf("%d session(s) exist with the generated session_secret", existing)
	}

	cookie := http.Cookie{
		Name:     "moneygo-session",
		Value:    session_secret,
		Path:     "/",
		Domain:   r.URL.Host,
		Expires:  time.Now().AddDate(0, 1, 0), // a month from now
		Secure:   true,
		HttpOnly: true,
	}

	s.SessionSecret = session_secret
	s.UserId = userid
	s.Created = time.Now()
	s.Expires = cookie.Expires

	err = tx.Insert(&s)
	if err != nil {
		return nil, err
	}
	return &NewSessionWriter{&s, &cookie}, nil
}

func SessionHandler(r *http.Request, context *Context) ResponseWriterWriter {
	if r.Method == "POST" || r.Method == "PUT" {
		var user User
		if err := ReadJSON(r, &user); err != nil {
			return NewError(3 /*Invalid Request*/)
		}

		dbuser, err := GetUserByUsername(context.Tx, user.Username)
		if err != nil {
			return NewError(2 /*Unauthorized Access*/)
		}

		user.HashPassword()
		if user.PasswordHash != dbuser.PasswordHash {
			return NewError(2 /*Unauthorized Access*/)
		}

		err = DeleteSessionIfExists(context.Tx, r)
		if err != nil {
			log.Print(err)
			return NewError(999 /*Internal Error*/)
		}

		sessionwriter, err := NewSession(context.Tx, r, dbuser.UserId)
		if err != nil {
			log.Print(err)
			return NewError(999 /*Internal Error*/)
		}
		return sessionwriter
	} else if r.Method == "GET" {
		s, err := GetSession(context.Tx, r)
		if err != nil {
			return NewError(1 /*Not Signed In*/)
		}

		return s
	} else if r.Method == "DELETE" {
		err := DeleteSessionIfExists(context.Tx, r)
		if err != nil {
			log.Print(err)
			return NewError(999 /*Internal Error*/)
		}
		return SuccessWriter{}
	}
	return NewError(3 /*Invalid Request*/)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Ceph-CSI Authors.

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

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

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

package liveness

import (
	"context"
	"time"

	"github.com/ceph/ceph-csi/internal/util"

	connlib "github.com/kubernetes-csi/csi-lib-utils/connection"
	"github.com/kubernetes-csi/csi-lib-utils/metrics"
	"github.com/kubernetes-csi/csi-lib-utils/rpc"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/grpc"
	klog "k8s.io/klog/v2"
)

var (
	liveness = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "csi",
		Name:      "liveness",
		Help:      "Liveness Probe",
	})
)

func getLiveness(timeout time.Duration, csiConn *grpc.ClientConn) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	util.TraceLogMsg("Sending probe request to CSI driver")
	ready, err := rpc.Probe(ctx, csiConn)
	if err != nil {
		liveness.Set(0)
		klog.Errorf("health check failed: %v", err)
		return
	}

	if !ready {
		liveness.Set(0)
		klog.Error("driver responded but is not ready")
		return
	}
	liveness.Set(1)
	util.ExtendedLogMsg("Health check succeeded")
}

func recordLiveness(endpoint, drivername string, pollTime, timeout time.Duration) {
	liveMetricsManager := metrics.NewCSIMetricsManager(drivername)
	// register prometheus metrics
	err := prometheus.Register(liveness)
	if err != nil {
		klog.Fatalln(err)
	}

	csiConn, err := connlib.Connect(endpoint, liveMetricsManager)
	if err != nil {
		// connlib should retry forever so a returned error should mean
		// the grpc client is misconfigured rather than an error on the network
		klog.Fatalf("failed to establish connection to CSI driver: %v", err)
	}

	// get liveness periodically
	ticker := time.NewTicker(pollTime)
	defer ticker.Stop()
	for range ticker.C {
		getLiveness(timeout, csiConn)
	}
}

// Run starts liveness collection and prometheus endpoint.
func Run(conf *util.Config) {
	util.ExtendedLogMsg("Liveness Running")

	// start liveness collection
	go recordLiveness(conf.Endpoint, conf.DriverName, conf.PollTime, conf.PoolTimeout)

	// start up prometheus endpoint
	util.StartMetricsServer(conf)
}
<commit_msg>util: replace klog with util logger in liveness.go<commit_after>/*
Copyright 2019 The Ceph-CSI Authors.

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

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

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

package liveness

import (
	"context"
	"time"

	"github.com/ceph/ceph-csi/internal/util"

	connlib "github.com/kubernetes-csi/csi-lib-utils/connection"
	"github.com/kubernetes-csi/csi-lib-utils/metrics"
	"github.com/kubernetes-csi/csi-lib-utils/rpc"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/grpc"
)

var (
	liveness = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "csi",
		Name:      "liveness",
		Help:      "Liveness Probe",
	})
)

func getLiveness(timeout time.Duration, csiConn *grpc.ClientConn) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	util.TraceLogMsg("Sending probe request to CSI driver")
	ready, err := rpc.Probe(ctx, csiConn)
	if err != nil {
		liveness.Set(0)
		util.ErrorLogMsg("health check failed: %v", err)
		return
	}

	if !ready {
		liveness.Set(0)
		util.ErrorLogMsg("driver responded but is not ready")
		return
	}
	liveness.Set(1)
	util.ExtendedLogMsg("Health check succeeded")
}

func recordLiveness(endpoint, drivername string, pollTime, timeout time.Duration) {
	liveMetricsManager := metrics.NewCSIMetricsManager(drivername)
	// register prometheus metrics
	err := prometheus.Register(liveness)
	if err != nil {
		util.FatalLogMsg(err.Error())
	}

	csiConn, err := connlib.Connect(endpoint, liveMetricsManager)
	if err != nil {
		// connlib should retry forever so a returned error should mean
		// the grpc client is misconfigured rather than an error on the network
		util.FatalLogMsg("failed to establish connection to CSI driver: %v", err)
	}

	// get liveness periodically
	ticker := time.NewTicker(pollTime)
	defer ticker.Stop()
	for range ticker.C {
		getLiveness(timeout, csiConn)
	}
}

// Run starts liveness collection and prometheus endpoint.
func Run(conf *util.Config) {
	util.ExtendedLogMsg("Liveness Running")

	// start liveness collection
	go recordLiveness(conf.Endpoint, conf.DriverName, conf.PollTime, conf.PoolTimeout)

	// start up prometheus endpoint
	util.StartMetricsServer(conf)
}
<|endoftext|>


<commit_before>package suites

import (
	"fmt"
	"os/exec"
	"strings"
	"time"

	"github.com/authelia/authelia/v4/internal/utils"
)

var kindImageName = "authelia-kind-proxy"
var dockerCmdLine = fmt.Sprintf("docker-compose -p authelia -f internal/suites/docker-compose.yml -f internal/suites/example/compose/kind/docker-compose.yml run --rm %s", kindImageName)

// Kind used for running kind commands.
type Kind struct{}

func kindCommand(cmdline string) *exec.Cmd {
	cmd := fmt.Sprintf("%s %s", dockerCmdLine, cmdline)
	return utils.Shell(cmd)
}

// CreateCluster create a new Kubernetes cluster.
func (k Kind) CreateCluster() error {
	cmd := kindCommand("kind create cluster --config /etc/kind/config.yml")
	if err := cmd.Run(); err != nil {
		return err
	}

	cmd = kindCommand("patch-kubeconfig.sh")
	if err := cmd.Run(); err != nil {
		return err
	}

	// This command is necessary to fix the coredns loop detected when using user-defined docker network.
	// In that case /etc/resolv.conf use 127.0.0.11 as DNS and CoreDNS thinks it is talking to itself which is wrong.
	// This IP is the docker internal DNS so it is safe to disable the loop check.
	cmd = kindCommand("sh -c 'kubectl -n kube-system get configmap/coredns -o yaml | grep -v loop | kubectl replace -f -'")
	err := cmd.Run()

	return err
}

// DeleteCluster delete a Kubernetes cluster.
func (k Kind) DeleteCluster() error {
	cmd := kindCommand("kind delete cluster")
	return cmd.Run()
}

// ClusterExists check whether a cluster exists.
func (k Kind) ClusterExists() (bool, error) {
	cmd := kindCommand("kind get clusters")
	cmd.Stdout = nil
	cmd.Stderr = nil
	output, err := cmd.Output()

	if err != nil {
		return false, err
	}

	return strings.Contains(string(output), "kind"), nil
}

// LoadImage load an image in the Kubernetes container.
func (k Kind) LoadImage(imageName string) error {
	cmd := kindCommand(fmt.Sprintf("kind load docker-image %s", imageName))
	return cmd.Run()
}

// Kubectl used for running kubectl commands.
type Kubectl struct{}

// StartProxy start a proxy.
func (k Kubectl) StartProxy() error {
	cmd := utils.Shell("docker-compose -p authelia -f internal/suites/docker-compose.yml -f internal/suites/example/compose/kind/docker-compose.yml up -d authelia-kind-proxy")
	return cmd.Run()
}

// StopProxy stop a proxy.
func (k Kubectl) StopProxy() error {
	cmd := utils.Shell("docker-compose -p authelia -f internal/suites/docker-compose.yml -f internal/suites/example/compose/kind/docker-compose.yml rm -s -f authelia-kind-proxy")
	return cmd.Run()
}

// StartDashboard start Kube dashboard.
func (k Kubectl) StartDashboard() error {
	if err := kindCommand("sh -c 'cd /authelia && ./bootstrap-dashboard.sh'").Run(); err != nil {
		return err
	}

	err := utils.Shell("docker-compose -p authelia -f internal/suites/docker-compose.yml -f internal/suites/example/compose/kind/docker-compose.yml up -d kube-dashboard").Run()

	return err
}

// StopDashboard stop kube dashboard.
func (k Kubectl) StopDashboard() error {
	cmd := utils.Shell("docker-compose -p authelia -f internal/suites/docker-compose.yml -f internal/suites/example/compose/kind/docker-compose.yml rm -s -f kube-dashboard")
	return cmd.Run()
}

// DeployThirdparties deploy thirdparty services (ldap, db, ingress controllers, etc...).
func (k Kubectl) DeployThirdparties() error {
	cmd := kindCommand("sh -c 'cd /authelia && ./bootstrap.sh'")
	return cmd.Run()
}

// DeployAuthelia deploy Authelia application.
func (k Kubectl) DeployAuthelia() error {
	cmd := kindCommand("sh -c 'cd /authelia && ./bootstrap-authelia.sh'")
	return cmd.Run()
}

// WaitPodsReady wait for all pods to be ready.
func (k Kubectl) WaitPodsReady(timeout time.Duration) error {
	return utils.CheckUntil(5*time.Second, timeout, func() (bool, error) {
		cmd := kindCommand("kubectl get -n authelia pods --no-headers")
		cmd.Stdout = nil
		cmd.Stderr = nil
		output, _ := cmd.Output()

		lines := strings.Split(string(output), "\n")

		nonEmptyLines := make([]string, 0)
		for _, line := range lines {
			if line != "" {
				nonEmptyLines = append(nonEmptyLines, line)
			}
		}

		for _, line := range nonEmptyLines {
			if !strings.Contains(line, "1/1") {
				return false, nil
			}
		}
		return true, nil
	})
}
<commit_msg>ci: fix docker-compose tty issue (#3496)<commit_after>package suites

import (
	"fmt"
	"os/exec"
	"strings"
	"time"

	"github.com/authelia/authelia/v4/internal/utils"
)

var kindImageName = "authelia-kind-proxy"
var dockerCmdLine = fmt.Sprintf("docker-compose -p authelia -f internal/suites/docker-compose.yml -f internal/suites/example/compose/kind/docker-compose.yml run -T --rm %s", kindImageName)

// Kind used for running kind commands.
type Kind struct{}

func kindCommand(cmdline string) *exec.Cmd {
	cmd := fmt.Sprintf("%s %s", dockerCmdLine, cmdline)
	return utils.Shell(cmd)
}

// CreateCluster create a new Kubernetes cluster.
func (k Kind) CreateCluster() error {
	cmd := kindCommand("kind create cluster --config /etc/kind/config.yml")
	if err := cmd.Run(); err != nil {
		return err
	}

	cmd = kindCommand("patch-kubeconfig.sh")
	if err := cmd.Run(); err != nil {
		return err
	}

	// This command is necessary to fix the coredns loop detected when using user-defined docker network.
	// In that case /etc/resolv.conf use 127.0.0.11 as DNS and CoreDNS thinks it is talking to itself which is wrong.
	// This IP is the docker internal DNS so it is safe to disable the loop check.
	cmd = kindCommand("sh -c 'kubectl -n kube-system get configmap/coredns -o yaml | grep -v loop | kubectl replace -f -'")
	err := cmd.Run()

	return err
}

// DeleteCluster delete a Kubernetes cluster.
func (k Kind) DeleteCluster() error {
	cmd := kindCommand("kind delete cluster")
	return cmd.Run()
}

// ClusterExists check whether a cluster exists.
func (k Kind) ClusterExists() (bool, error) {
	cmd := kindCommand("kind get clusters")
	cmd.Stdout = nil
	cmd.Stderr = nil
	output, err := cmd.Output()

	if err != nil {
		return false, err
	}

	return strings.Contains(string(output), "kind"), nil
}

// LoadImage load an image in the Kubernetes container.
func (k Kind) LoadImage(imageName string) error {
	cmd := kindCommand(fmt.Sprintf("kind load docker-image %s", imageName))
	return cmd.Run()
}

// Kubectl used for running kubectl commands.
type Kubectl struct{}

// StartProxy start a proxy.
func (k Kubectl) StartProxy() error {
	cmd := utils.Shell("docker-compose -p authelia -f internal/suites/docker-compose.yml -f internal/suites/example/compose/kind/docker-compose.yml up -d authelia-kind-proxy")
	return cmd.Run()
}

// StopProxy stop a proxy.
func (k Kubectl) StopProxy() error {
	cmd := utils.Shell("docker-compose -p authelia -f internal/suites/docker-compose.yml -f internal/suites/example/compose/kind/docker-compose.yml rm -s -f authelia-kind-proxy")
	return cmd.Run()
}

// StartDashboard start Kube dashboard.
func (k Kubectl) StartDashboard() error {
	if err := kindCommand("sh -c 'cd /authelia && ./bootstrap-dashboard.sh'").Run(); err != nil {
		return err
	}

	err := utils.Shell("docker-compose -p authelia -f internal/suites/docker-compose.yml -f internal/suites/example/compose/kind/docker-compose.yml up -d kube-dashboard").Run()

	return err
}

// StopDashboard stop kube dashboard.
func (k Kubectl) StopDashboard() error {
	cmd := utils.Shell("docker-compose -p authelia -f internal/suites/docker-compose.yml -f internal/suites/example/compose/kind/docker-compose.yml rm -s -f kube-dashboard")
	return cmd.Run()
}

// DeployThirdparties deploy thirdparty services (ldap, db, ingress controllers, etc...).
func (k Kubectl) DeployThirdparties() error {
	cmd := kindCommand("sh -c 'cd /authelia && ./bootstrap.sh'")
	return cmd.Run()
}

// DeployAuthelia deploy Authelia application.
func (k Kubectl) DeployAuthelia() error {
	cmd := kindCommand("sh -c 'cd /authelia && ./bootstrap-authelia.sh'")
	return cmd.Run()
}

// WaitPodsReady wait for all pods to be ready.
func (k Kubectl) WaitPodsReady(timeout time.Duration) error {
	return utils.CheckUntil(5*time.Second, timeout, func() (bool, error) {
		cmd := kindCommand("kubectl get -n authelia pods --no-headers")
		cmd.Stdout = nil
		cmd.Stderr = nil
		output, _ := cmd.Output()

		lines := strings.Split(string(output), "\n")

		nonEmptyLines := make([]string, 0)
		for _, line := range lines {
			if line != "" {
				nonEmptyLines = append(nonEmptyLines, line)
			}
		}

		for _, line := range nonEmptyLines {
			if !strings.Contains(line, "1/1") {
				return false, nil
			}
		}
		return true, nil
	})
}
<|endoftext|>


<commit_before>/*
 *
 * Copyright 2018 gRPC 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.
 *
 */

// This binary can only run on Google Cloud Platform (GCP).
package main

import (
	"flag"
	"net"
	"strings"

	grpc "google.golang.org/grpc"
	"google.golang.org/grpc/credentials/alts"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/interop"
	testpb "google.golang.org/grpc/interop/grpc_testing"
)

const (
	udsAddrPrefix = "unix:"
)

var (
	hsAddr     = flag.String("alts_handshaker_service_address", "", "ALTS handshaker gRPC service address")
	serverAddr = flag.String("server_address", ":8080", "The address on which the server is listening. Only two types of addresses are supported, 'host:port' and 'unix:/path'.")
)

func main() {
	flag.Parse()

	// If the server address starts with `unix:`, then we have a UDS address.
	network := "tcp"
	address := *serverAddr
	if strings.HasPrefix(address, udsAddrPrefix) {
		network = "unix"
		address = strings.TrimPrefix(address, udsAddrPrefix)
	}
	lis, err := net.Listen(network, address)
	if err != nil {
		grpclog.Fatalf("gRPC Server: failed to start the server at %v: %v", address, err)
	}
	opts := alts.DefaultServerOptions()
	if *hsAddr != "" {
		opts.HandshakerServiceAddress = *hsAddr
	}
	altsTC := alts.NewServerCreds(opts)
	grpcServer := grpc.NewServer(grpc.Creds(altsTC))
	testpb.RegisterTestServiceServer(grpcServer, interop.NewTestServer())
	grpcServer.Serve(lis)
}
<commit_msg>credentials/alts: Add example of authz in ALTS (#2814)<commit_after>/*
 *
 * Copyright 2018 gRPC 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.
 *
 */

// This binary can only run on Google Cloud Platform (GCP).
package main

import (
	"context"
	"flag"
	"net"
	"strings"

	grpc "google.golang.org/grpc"
	"google.golang.org/grpc/credentials/alts"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/interop"
	testpb "google.golang.org/grpc/interop/grpc_testing"
	"google.golang.org/grpc/tap"
)

const (
	udsAddrPrefix = "unix:"
)

var (
	hsAddr     = flag.String("alts_handshaker_service_address", "", "ALTS handshaker gRPC service address")
	serverAddr = flag.String("server_address", ":8080", "The address on which the server is listening. Only two types of addresses are supported, 'host:port' and 'unix:/path'.")
)

func main() {
	flag.Parse()

	// If the server address starts with `unix:`, then we have a UDS address.
	network := "tcp"
	address := *serverAddr
	if strings.HasPrefix(address, udsAddrPrefix) {
		network = "unix"
		address = strings.TrimPrefix(address, udsAddrPrefix)
	}
	lis, err := net.Listen(network, address)
	if err != nil {
		grpclog.Fatalf("gRPC Server: failed to start the server at %v: %v", address, err)
	}
	opts := alts.DefaultServerOptions()
	if *hsAddr != "" {
		opts.HandshakerServiceAddress = *hsAddr
	}
	altsTC := alts.NewServerCreds(opts)
	grpcServer := grpc.NewServer(grpc.Creds(altsTC), grpc.InTapHandle(authz))
	testpb.RegisterTestServiceServer(grpcServer, interop.NewTestServer())
	grpcServer.Serve(lis)
}

// authz shows how to access client information at the server side to perform
// application-layer authorization checks.
func authz(ctx context.Context, info *tap.Info) (context.Context, error) {
	authInfo, err := alts.AuthInfoFromContext(ctx)
	if err != nil {
		return nil, err
	}
	// Access all alts.AuthInfo data:
	grpclog.Infof("authInfo.ApplicationProtocol() = %v", authInfo.ApplicationProtocol())
	grpclog.Infof("authInfo.RecordProtocol() = %v", authInfo.RecordProtocol())
	grpclog.Infof("authInfo.SecurityLevel() = %v", authInfo.SecurityLevel())
	grpclog.Infof("authInfo.PeerServiceAccount() = %v", authInfo.PeerServiceAccount())
	grpclog.Infof("authInfo.LocalServiceAccount() = %v", authInfo.LocalServiceAccount())
	grpclog.Infof("authInfo.PeerRPCVersions() = %v", authInfo.PeerRPCVersions())
	grpclog.Infof("info.FullMethodName = %v", info.FullMethodName)
	return ctx, nil
}
<|endoftext|>


<commit_before>package carchivum

import (
	"archive/zip"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/MichaelTJones/walk"
)

// Zip handles .zip archives.
type Zip struct {
	Car
	writer  *zip.Writer
	fwriter *os.File
}

// NewZip returns an initialized Zip struct ready for use.
func NewZip() *Zip {
	return &Zip{
		Car: Car{t0: time.Now()},
	}
}

// Create creates a zip file from src in the dst
func (z *Zip) Create(dst string, src ...string) (cnt int, err error) {
	// If there isn't a destination, return err
	if dst == "" {
		err = fmt.Errorf("destination required to create a zip archive")
		log.Print(err)
		return 0, err
	}

	// If there aren't any sources, return err
	if len(dst) == 0 {
		err = fmt.Errorf("a source is required to create a zip archive")
		log.Print(err)
		return 0, err
	}

	// See if we can create the destination file before processing
	z.fwriter, err = os.OpenFile(dst, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		log.Print(err)
		return 0, err
	}
	defer z.fwriter.Close()

	buf := new(bytes.Buffer)
	z.writer = zip.NewWriter(buf)
	defer z.writer.Close()

	// Set up the file queue and its drain.
	z.FileCh = make(chan *os.File)
	wait, err := z.write()
	if err != nil {
		log.Print(err)
		return 0, err
	}

	var fullPath string
	// Walk the sources, add each file to the queue.
	// This isn't limited as a large number of sources is not expected.
	//
	visitor := func(p string, fi os.FileInfo, err error) error {
		return z.AddFile(fullPath, p, fi, err)
	}

	var wg sync.WaitGroup
	wg.Add(len(src) - 1)
	for _, source := range src {
		// first get the absolute, its needed either way
		fullPath, err = filepath.Abs(source)
		if err != nil {
			log.Print(err)
			return 0, err
		}

		err = walk.Walk(fullPath, visitor)
		if err != nil {
			log.Print(err)
			return 0, err
		}
	}

	wg.Wait()

	close(z.FileCh)
	wait.Wait()

	z.writer.Close()

	// Copy the zip
	_, err = z.fwriter.Write(buf.Bytes())
	if err != nil {
		log.Print(err)
		return 0, err
	}

	z.fwriter.Close()
	z.setDelta()
	return int(z.Car.files), nil
}

// ZipBytes takes a string and bytes and returns a zip archive of the bytes
// using the name.
func ZipBytes(b []byte, name string) (n int, zipped []byte, err error) {
	buf := new(bytes.Buffer)
	w := zip.NewWriter(buf)
	defer w.Close() // defer for convenience, though it may already be closed
	f, err := w.Create(name)
	if err != nil {
		log.Print(err)
		return 0, zipped, err
	}

	n, err = f.Write(b)
	if err != nil {
		log.Print(err)
		return n, zipped, err
	}
	w.Close() // we need to close it to get the bytes.
	return n, buf.Bytes(), err
}

func copyTo(w io.Writer, z *zip.File) (int64, error) {
	f, err := z.Open()
	if err != nil {
		return 0, err
	}
	defer f.Close()

	return io.Copy(w, f)
}

//
// Because zip can't be parallized because  `Create/CreateHEader` implicitly
// closes the writer and I don't feel like writing a parallized zip writer,
// we spawn a new goroutine for each file to read and pipe them to the zipper
// goroutine.
//
// SEE where to add defer
func (z *Zip) write() (*sync.WaitGroup, error) {
	var wg sync.WaitGroup
	wg.Add(1)
	go func() error {
		defer wg.Done()

		for f := range z.FileCh {
			defer f.Close()
			info, err := f.Stat()
			if err != nil {
				log.Print(err)
				return err
			}

			if info.IsDir() {
				continue
			}

			header, err := zip.FileInfoHeader(info)
			if err != nil {
				log.Print(err)
				return err
			}

			header.Name = f.Name()

			fw, err := z.writer.CreateHeader(header)
			if err != nil {
				log.Print(err)
				return err
			}

			_, err = io.Copy(fw, f)
			if err != nil {
				log.Print(err)
				return err
			}

		}

		return nil
	}()

	return &wg, nil
}

// Extract the content of src, a zip archive, to dst.
func (z *Zip) Extract(dst, src string) error {
	r, err := zip.OpenReader(src)
	if err != nil {
		log.Print(err)
		return err
	}
	defer r.Close()

	for _, f := range r.File {
		rc, err := f.Open()
		if err != nil {
			log.Print(err)
			return err
		}
		err = os.MkdirAll(filepath.Join(dst, filepath.Dir(f.Name)), 0755)
		if err != nil {
			log.Print(err)
			return err
		}
		dF, err := os.Create(filepath.Join(dst, f.Name))
		if err != nil {
			log.Print(err)
			return err
		}
		_, err = io.Copy(dF, rc)
		if err != nil {
			log.Print(err)
			return err
		}
		rc.Close()
		dF.Close()
	}
	return nil
}
<commit_msg>Zip.writer -> Zip.Writer && Zip.fwriter -> Zip.File<commit_after>package carchivum

import (
	"archive/zip"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/MichaelTJones/walk"
)

// Zip handles .zip archives.
type Zip struct {
	Car
	*zip.Writer
	*os.File
}

// NewZip returns an initialized Zip struct ready for use.
func NewZip() *Zip {
	return &Zip{
		Car: Car{t0: time.Now()},
	}
}

// Create creates a zip file from src in the dst
func (z *Zip) Create(dst string, src ...string) (cnt int, err error) {
	// If there isn't a destination, return err
	if dst == "" {
		err = fmt.Errorf("destination required to create a zip archive")
		log.Print(err)
		return 0, err
	}

	// If there aren't any sources, return err
	if len(dst) == 0 {
		err = fmt.Errorf("a source is required to create a zip archive")
		log.Print(err)
		return 0, err
	}

	// See if we can create the destination file before processing
	z.File, err = os.OpenFile(dst, os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		log.Print(err)
		return 0, err
	}
	defer z.File.Close()

	buf := new(bytes.Buffer)
	z.Writer = zip.NewWriter(buf)
	defer z.Writer.Close()

	// Set up the file queue and its drain.
	z.FileCh = make(chan *os.File)
	wait, err := z.write()
	if err != nil {
		log.Print(err)
		return 0, err
	}

	var fullPath string
	// Walk the sources, add each file to the queue.
	// This isn't limited as a large number of sources is not expected.
	//
	visitor := func(p string, fi os.FileInfo, err error) error {
		return z.AddFile(fullPath, p, fi, err)
	}

	var wg sync.WaitGroup
	wg.Add(len(src) - 1)
	for _, source := range src {
		// first get the absolute, its needed either way
		fullPath, err = filepath.Abs(source)
		if err != nil {
			log.Print(err)
			return 0, err
		}

		err = walk.Walk(fullPath, visitor)
		if err != nil {
			log.Print(err)
			return 0, err
		}
	}

	wg.Wait()

	close(z.FileCh)
	wait.Wait()

	z.Writer.Close()

	// Copy the zip
	_, err = z.File.Write(buf.Bytes())
	if err != nil {
		log.Print(err)
		return 0, err
	}

	z.File.Close()
	z.setDelta()
	return int(z.Car.files), nil
}

// ZipBytes takes a string and bytes and returns a zip archive of the bytes
// using the name.
func ZipBytes(b []byte, name string) (n int, zipped []byte, err error) {
	buf := new(bytes.Buffer)
	w := zip.NewWriter(buf)
	defer w.Close() // defer for convenience, though it may already be closed
	f, err := w.Create(name)
	if err != nil {
		log.Print(err)
		return 0, zipped, err
	}

	n, err = f.Write(b)
	if err != nil {
		log.Print(err)
		return n, zipped, err
	}
	w.Close() // we need to close it to get the bytes.
	return n, buf.Bytes(), err
}

func copyTo(w io.Writer, z *zip.File) (int64, error) {
	f, err := z.Open()
	if err != nil {
		return 0, err
	}
	defer f.Close()

	return io.Copy(w, f)
}

//
// Because zip can't be parallized because  `Create/CreateHEader` implicitly
// closes the writer and I don't feel like writing a parallized zip writer,
// we spawn a new goroutine for each file to read and pipe them to the zipper
// goroutine.
//
// SEE where to add defer
func (z *Zip) write() (*sync.WaitGroup, error) {
	var wg sync.WaitGroup
	wg.Add(1)
	go func() error {
		defer wg.Done()

		for f := range z.FileCh {
			defer f.Close()
			info, err := f.Stat()
			if err != nil {
				log.Print(err)
				return err
			}

			if info.IsDir() {
				continue
			}

			header, err := zip.FileInfoHeader(info)
			if err != nil {
				log.Print(err)
				return err
			}

			header.Name = f.Name()

			fw, err := z.Writer.CreateHeader(header)
			if err != nil {
				log.Print(err)
				return err
			}

			_, err = io.Copy(fw, f)
			if err != nil {
				log.Print(err)
				return err
			}

		}

		return nil
	}()

	return &wg, nil
}

// Extract the content of src, a zip archive, to dst.
func (z *Zip) Extract(dst, src string) error {
	r, err := zip.OpenReader(src)
	if err != nil {
		log.Print(err)
		return err
	}
	defer r.Close()

	for _, f := range r.File {
		rc, err := f.Open()
		if err != nil {
			log.Print(err)
			return err
		}
		err = os.MkdirAll(filepath.Join(dst, filepath.Dir(f.Name)), 0755)
		if err != nil {
			log.Print(err)
			return err
		}
		dF, err := os.Create(filepath.Join(dst, f.Name))
		if err != nil {
			log.Print(err)
			return err
		}
		_, err = io.Copy(dF, rc)
		if err != nil {
			log.Print(err)
			return err
		}
		rc.Close()
		dF.Close()
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package discoverspaces

import (
	"github.com/juju/errors"
	"github.com/juju/juju/api/discoverspaces"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/worker"
	"github.com/juju/loggo"
	"github.com/juju/names"
	"launchpad.net/tomb"
)

var logger = loggo.GetLogger("juju.discoverspaces")

type discoverspacesWorker struct {
	api      *discoverspaces.API
	tomb     tomb.Tomb
	observer *worker.EnvironObserver
}

// NewWorker returns a worker
func NewWorker(api *discoverspaces.API) worker.Worker {
	dw := &discoverspacesWorker{
		api: api,
	}
	go func() {
		defer dw.tomb.Done()
		dw.tomb.Kill(dw.loop())
	}()
	return dw
}

func (dw *discoverspacesWorker) Kill() {
	dw.tomb.Kill(nil)
}

func (dw *discoverspacesWorker) Wait() error {
	return dw.tomb.Wait()
}

func (dw *discoverspacesWorker) loop() (err error) {
	dw.observer, err = worker.NewEnvironObserver(dw.api)
	if err != nil {
		return err
	}
	defer func() {
		obsErr := worker.Stop(dw.observer)
		if err == nil {
			err = obsErr
		}
	}()
	environ := dw.observer.Environ()
	networkingEnviron, ok := environs.SupportsNetworking(environ)

	if ok {
		err = dw.handleSubnets(networkingEnviron)
		if err != nil {
			return errors.Trace(err)
		}
	}

	// TODO(mfoord): we'll have a watcher here checking if we need to
	// update the spaces/subnets definition.
	for {
	}
	return err
}

func (dw *discoverspacesWorker) handleSubnets(env environs.NetworkingEnviron) error {
	ok, err := env.SupportsSpaceDiscovery()
	if err != nil {
		return errors.Trace(err)
	}
	if !ok {
		// Nothing to do.
		return nil
	}
	providerSpaces, err := env.Spaces()
	if err != nil {
		return errors.Trace(err)
	}
	listSpacesResult, err := dw.api.ListSpaces()
	if err != nil {
		return errors.Trace(err)
	}
	stateSpaceMap := make(map[string]params.ProviderSpace)
	for _, space := range listSpacesResult.Results {
		stateSpaceMap[space.ProviderId] = space
	}

	// TODO(mfoord): we also need to attempt to delete spaces that no
	// longer exist, so long as they're not in use.
	for _, space := range providerSpaces {
		_, ok := stateSpaceMap[space.Name]
		if !ok {
			// We need to create the space.
			// XXX in the apiserver the name should be generated and
			// IsPublic set to false.
			args := params.CreateSpacesParams{}
			_, err := dw.api.CreateSpaces(args)
			if err != nil {
				return errors.Trace(err)
			}
		}
		// TODO(mfoord): currently no way of removing subnets, or
		// changing the space they're in, so we can only add ones we
		// don't already know about.
		for _, subnet := range space.Subnets {
			spaceTag, err := names.ParseSpaceTag(space.Name)
			if err != nil {
				return errors.Trace(err)
			}
			subnetTag, err := names.ParseSubnetTag(subnet.CIDR)
			if err != nil {
				return errors.Trace(err)
			}

			args := params.AddSubnetsParams{
				Subnets: []params.AddSubnetParams{{
					SubnetTag:        subnetTag.String(),
					SubnetProviderId: string(subnet.ProviderId),
					SpaceTag:         spaceTag.String(),
					Zones:            subnet.AvailabilityZones,
				}}}
			_, err = dw.api.AddSubnets(args)
			if err != nil {
				return errors.Trace(err)
			}
		}
	}
	return nil
}
<commit_msg>Listen for dying<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package discoverspaces

import (
	"github.com/juju/errors"
	"github.com/juju/juju/api/discoverspaces"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/worker"
	"github.com/juju/loggo"
	"github.com/juju/names"
	"launchpad.net/tomb"
)

var logger = loggo.GetLogger("juju.discoverspaces")

type discoverspacesWorker struct {
	api      *discoverspaces.API
	tomb     tomb.Tomb
	observer *worker.EnvironObserver
}

// NewWorker returns a worker
func NewWorker(api *discoverspaces.API) worker.Worker {
	dw := &discoverspacesWorker{
		api: api,
	}
	go func() {
		defer dw.tomb.Done()
		dw.tomb.Kill(dw.loop())
	}()
	return dw
}

func (dw *discoverspacesWorker) Kill() {
	dw.tomb.Kill(nil)
}

func (dw *discoverspacesWorker) Wait() error {
	return dw.tomb.Wait()
}

func (dw *discoverspacesWorker) loop() (err error) {
	dw.observer, err = worker.NewEnvironObserver(dw.api)
	if err != nil {
		return err
	}
	defer func() {
		obsErr := worker.Stop(dw.observer)
		if err == nil {
			err = obsErr
		}
	}()
	environ := dw.observer.Environ()
	networkingEnviron, ok := environs.SupportsNetworking(environ)

	if ok {
		err = dw.handleSubnets(networkingEnviron)
		if err != nil {
			return errors.Trace(err)
		}
	}

	// TODO(mfoord): we'll have a watcher here checking if we need to
	// update the spaces/subnets definition.
	dying := u.tomb.Dying()
	for {
		select {
		case <-dying:
			return nil
		}
	}
	return nil
}

func (dw *discoverspacesWorker) handleSubnets(env environs.NetworkingEnviron) error {
	ok, err := env.SupportsSpaceDiscovery()
	if err != nil {
		return errors.Trace(err)
	}
	if !ok {
		// Nothing to do.
		return nil
	}
	providerSpaces, err := env.Spaces()
	if err != nil {
		return errors.Trace(err)
	}
	listSpacesResult, err := dw.api.ListSpaces()
	if err != nil {
		return errors.Trace(err)
	}
	stateSpaceMap := make(map[string]params.ProviderSpace)
	for _, space := range listSpacesResult.Results {
		stateSpaceMap[space.ProviderId] = space
	}

	// TODO(mfoord): we also need to attempt to delete spaces and subnets
	// that no longer exist, so long as they're not in use.
	for _, space := range providerSpaces {
		_, ok := stateSpaceMap[space.Name]
		if !ok {
			// We need to create the space.
			// XXX in the apiserver the name should be generated and
			// IsPublic set to false.
			args := params.CreateSpacesParams{}
			_, err := dw.api.CreateSpaces(args)
			if err != nil {
				return errors.Trace(err)
			}
		}
		// TODO(mfoord): currently no way of removing subnets, or
		// changing the space they're in, so we can only add ones we
		// don't already know about.
		for _, subnet := range space.Subnets {
			spaceTag, err := names.ParseSpaceTag(space.Name)
			if err != nil {
				return errors.Trace(err)
			}
			subnetTag, err := names.ParseSubnetTag(subnet.CIDR)
			if err != nil {
				return errors.Trace(err)
			}

			args := params.AddSubnetsParams{
				Subnets: []params.AddSubnetParams{{
					SubnetTag:        subnetTag.String(),
					SubnetProviderId: string(subnet.ProviderId),
					SpaceTag:         spaceTag.String(),
					Zones:            subnet.AvailabilityZones,
				}}}
			_, err = dw.api.AddSubnets(args)
			if err != nil {
				return errors.Trace(err)
			}
		}
	}
	return nil
}
<|endoftext|>