<commit_before>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"github.com/ViBiOh/dashboard/auth"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"regexp"
)

const ignoredByteLogSize = 8
const tailSize = `100`
const start = `start`
const stop = `stop`

var eventsDemand = regexp.MustCompile(`^events (.+)(.?)`)
var logsDemand = regexp.MustCompile(`^logs (.+) (.+)`)
var statsDemand = regexp.MustCompile(`^stats (.+) (.+)`)
var busWebsocketRequest = regexp.MustCompile(`bus`)
var logWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/logs`)
var statsWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/stats`)
var eventsWebsocketRequest = regexp.MustCompile(`events`)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return hostCheck.MatchString(r.Host)
	},
}

func readUntilClose(user *auth.User, ws *websocket.Conn, name string) bool {
	messageType, _, err := ws.ReadMessage()

	if messageType == websocket.CloseMessage {
		return true
	}

	if err != nil {
		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure) {
			log.Printf(`[%s] Error while reading from %s socket: %v`, user.Username, name, err)
		}

		return true
	}

	return false
}

func readContent(user *auth.User, ws *websocket.Conn, name string, done chan<- int, content chan<- []byte) {
	for {
		messageType, message, err := ws.ReadMessage()

		if messageType == websocket.CloseMessage {
			close(done)
			return
		}

		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure) {
				log.Printf(`[%s] Error while reading from %s socket: %v`, user.Username, name, err)
			}

			close(done)
			return
		}

		content <- message
	}
}

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, *auth.User, error) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	user, err := auth.IsAuthenticatedByAuth(string(basicAuth))
	if err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))

		defer ws.Close()
		return nil, nil, err
	}

	return ws, user, nil
}

func logsContainerWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	logs, err := docker.ContainerLogs(ctx, string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	if err != nil {
		log.Print(err)
		return
	}
	defer logs.Close()

	go func() {
		defer cancel()

		scanner := bufio.NewScanner(logs)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				return

			default:
				logLine := scanner.Bytes()
				if len(logLine) > ignoredByteLogSize {
					if err = ws.WriteMessage(websocket.TextMessage, logLine[ignoredByteLogSize:]); err != nil {
						log.Printf(`[%s] Error while writing to logs socket: %v`, user.Username, err)
						return
					}
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			break

		default:
			if readUntilClose(user, ws, `logs`) {
				return
			}
		}
	}
}

func eventsWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		return
	}
	defer ws.Close()

	filtersArgs := filters.NewArgs()
	if labelFilters(&filtersArgs, user, nil) != nil {
		log.Printf(`[%s] Error while defining label filters: %v`, user.Username, err)
		return
	}
	if eventFilters(&filtersArgs) != nil {
		log.Printf(`[%s] Error while defining event filters: %v`, user.Username, err)
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	messages, errors := docker.Events(ctx, types.EventsOptions{Filters: filtersArgs})

	go func() {
		defer cancel()

		for {
			select {
			case <-ctx.Done():
				return

			case message := <-messages:
				messageJSON, err := json.Marshal(message)
				if err != nil {
					log.Printf(`[%s] Error while marshalling event: %v`, user.Username, err)
					return
				}

				if err = ws.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
					log.Printf(`[%s] Error while writing to events socket: %v`, user.Username, err)
					return
				}

			case err := <-errors:
				log.Printf(`[%s] Error while reading events: %v`, user.Username, err)
				return
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			break

		default:
			if readUntilClose(user, ws, `events`) {
				return
			}
		}
	}
}

func statsWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	stats, err := docker.ContainerStats(ctx, string(containerID), true)
	if err != nil {
		log.Print(err)
		return
	}
	defer stats.Body.Close()

	go func() {
		defer cancel()

		scanner := bufio.NewScanner(stats.Body)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				log.Printf(`[%s] Stats context is over for writing`, user.Username)
				return

			default:
				if err = ws.WriteMessage(websocket.TextMessage, scanner.Bytes()); err != nil {
					log.Printf(`[%s] Error while writing to stats socket: %v`, user.Username, err)
					return
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			log.Printf(`[%s] Stats context is over for reading`, user.Username)
			return

		default:
			if readUntilClose(user, ws, `stats`) {
				return
			}
		}
	}
}

func streamLogs(ctx context.Context, cancel context.CancelFunc, user *auth.User, containerID string, output chan<- []byte) {
	logs, err := docker.ContainerLogs(ctx, containerID, types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	defer cancel()

	if err != nil {
		log.Printf(`[%s] Logs opening in error: %v`, user.Username, err)
		return
	}
	defer logs.Close()

	scanner := bufio.NewScanner(logs)
	log.Printf(`[%s] Logs streaming started for %s`, user.Username, containerID)

	for scanner.Scan() {
		logLine := scanner.Bytes()
		if len(logLine) > ignoredByteLogSize {
			output <- logLine[ignoredByteLogSize:]
		}
	}

	log.Printf(`[%s] Logs streaming ended for %s`, user.Username, containerID)
}

func streamEvents(ctx context.Context, cancel context.CancelFunc, user *auth.User, _ string, output chan<- []byte) {
	defer cancel()

	filtersArgs := filters.NewArgs()
	if err := labelFilters(&filtersArgs, user, nil); err != nil {
		log.Printf(`[%s] Events opening in error: %v`, user.Username, err)
		return
	}
	if err := eventFilters(&filtersArgs); err != nil {
		log.Printf(`[%s] Events opening in error: %v`, user.Username, err)
		return
	}

	messages, errors := docker.Events(ctx, types.EventsOptions{Filters: filtersArgs})

	log.Printf(`[%s] Events streaming started`, user.Username)
	for {
		select {
		case <-ctx.Done():
			log.Printf(`[%s] Events streaming ended`, user.Username)
			return

		case message := <-messages:
			messageJSON, err := json.Marshal(message)
			if err != nil {
				log.Printf(`[%s] Events marshalling in error: %v`, user.Username, err)
				cancel()
			} else {
				output <- messageJSON
			}

		case err := <-errors:
			log.Printf(`[%s] Events reading in error: %v`, user.Username, err)
			cancel()
		}
	}
}

func streamStats(ctx context.Context, cancel context.CancelFunc, user *auth.User, containerID string, output chan<- []byte) {
	stats, err := docker.ContainerStats(ctx, containerID, true)
	defer cancel()

	if err != nil {
		log.Printf(`[%s] Stats opening in error for %s: %v`, user.Username, containerID, err)
		return
	}
	defer stats.Body.Close()

	scanner := bufio.NewScanner(stats.Body)
	log.Printf(`[%s] Stats streaming started for %s`, user.Username, containerID)

	for scanner.Scan() {
		output <- scanner.Bytes()
	}

	log.Printf(`[%s] Stats streaming ended for %s`, user.Username, containerID)
}

func handleBusDemand(user *auth.User, name string, input []byte, demand *regexp.Regexp, cancel context.CancelFunc, output chan<- []byte, streamFn func(context.Context, context.CancelFunc, *auth.User, string, chan<- []byte)) context.CancelFunc {
	demandGroups := demand.FindSubmatch(input)
	if len(demandGroups) < 2 {
		log.Printf(`[%s] Unable to parse bus demand %s for %s`, user.Username, input, name)
	}

	containerID := string(demandGroups[1])
	action := string(demandGroups[2])

	if action == stop && cancel != nil {
		log.Printf(`[%s] Stopping %s stream`, user.Username, name)
		cancel()
	} else if action == start {
		log.Printf(`[%s] Starting %s stream for %s`, user.Username, name, containerID)

		if cancel != nil {
			log.Printf(`[%s] Cancelling previous %s stream`, user.Username, name)
			cancel()
		}

		ctx, newCancel := context.WithCancel(context.Background())
		go streamFn(ctx, newCancel, user, string(containerID), output)

		return newCancel
	}

	return nil
}

func busWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	done := make(chan int)

	output := make(chan []byte)
	defer close(output)

	input := make(chan []byte)
	defer close(input)

	go readContent(user, ws, `streaming`, done, input)
	log.Printf(`[%s] Streaming started`, user.Username)

	var eventsCancelFunc context.CancelFunc
	var logsCancelFunc context.CancelFunc
	var statsCancelFunc context.CancelFunc

	for {
		select {
		case <-done:
			log.Printf(`[%s] Streaming ended`, user.Username)
			return

		case inputBytes := <-input:
			if eventsDemand.Match(inputBytes) {
				eventsCancelFunc = handleBusDemand(user, `events`, inputBytes, eventsDemand, eventsCancelFunc, output, streamEvents)
				if eventsCancelFunc != nil {
					defer eventsCancelFunc()
				}
			} else if logsDemand.Match(inputBytes) {
				logsCancelFunc = handleBusDemand(user, `logs`, inputBytes, logsDemand, logsCancelFunc, output, streamLogs)
				if logsCancelFunc != nil {
					defer logsCancelFunc()
				}
			} else if statsDemand.Match(inputBytes) {
				statsCancelFunc = handleBusDemand(user, `stats`, inputBytes, statsDemand, statsCancelFunc, output, streamStats)
				if statsCancelFunc != nil {
					defer statsCancelFunc()
				}
			}

		case outputBytes := <-output:
			if err = ws.WriteMessage(websocket.TextMessage, outputBytes); err != nil {
				log.Printf(`[%s] Error while writing to streaming: %v`, user.Username, err)
				close(done)
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	urlPath := []byte(r.URL.Path)

	if logWebsocketRequest.Match(urlPath) {
		logsContainerWebsocketHandler(w, r, logWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if eventsWebsocketRequest.Match((urlPath)) {
		eventsWebsocketHandler(w, r)
	} else if statsWebsocketRequest.Match(urlPath) {
		statsWebsocketHandler(w, r, statsWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if busWebsocketRequest.Match(urlPath) {
		busWebsocketHandler(w, r)
	}
}
<commit_msg>Inverting group matching for case when no containerID<commit_after>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"github.com/ViBiOh/dashboard/auth"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"regexp"
)

const ignoredByteLogSize = 8
const tailSize = `100`
const start = `start`
const stop = `stop`

var eventsDemand = regexp.MustCompile(`^events (.+)`)
var logsDemand = regexp.MustCompile(`^logs (.+) (.+)`)
var statsDemand = regexp.MustCompile(`^stats (.+) (.+)`)
var busWebsocketRequest = regexp.MustCompile(`bus`)
var logWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/logs`)
var statsWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/stats`)
var eventsWebsocketRequest = regexp.MustCompile(`events`)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return hostCheck.MatchString(r.Host)
	},
}

func readUntilClose(user *auth.User, ws *websocket.Conn, name string) bool {
	messageType, _, err := ws.ReadMessage()

	if messageType == websocket.CloseMessage {
		return true
	}

	if err != nil {
		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure) {
			log.Printf(`[%s] Error while reading from %s socket: %v`, user.Username, name, err)
		}

		return true
	}

	return false
}

func readContent(user *auth.User, ws *websocket.Conn, name string, done chan<- int, content chan<- []byte) {
	for {
		messageType, message, err := ws.ReadMessage()

		if messageType == websocket.CloseMessage {
			close(done)
			return
		}

		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure) {
				log.Printf(`[%s] Error while reading from %s socket: %v`, user.Username, name, err)
			}

			close(done)
			return
		}

		content <- message
	}
}

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, *auth.User, error) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	user, err := auth.IsAuthenticatedByAuth(string(basicAuth))
	if err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))

		defer ws.Close()
		return nil, nil, err
	}

	return ws, user, nil
}

func logsContainerWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	logs, err := docker.ContainerLogs(ctx, string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	if err != nil {
		log.Print(err)
		return
	}
	defer logs.Close()

	go func() {
		defer cancel()

		scanner := bufio.NewScanner(logs)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				return

			default:
				logLine := scanner.Bytes()
				if len(logLine) > ignoredByteLogSize {
					if err = ws.WriteMessage(websocket.TextMessage, logLine[ignoredByteLogSize:]); err != nil {
						log.Printf(`[%s] Error while writing to logs socket: %v`, user.Username, err)
						return
					}
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			break

		default:
			if readUntilClose(user, ws, `logs`) {
				return
			}
		}
	}
}

func eventsWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		return
	}
	defer ws.Close()

	filtersArgs := filters.NewArgs()
	if labelFilters(&filtersArgs, user, nil) != nil {
		log.Printf(`[%s] Error while defining label filters: %v`, user.Username, err)
		return
	}
	if eventFilters(&filtersArgs) != nil {
		log.Printf(`[%s] Error while defining event filters: %v`, user.Username, err)
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	messages, errors := docker.Events(ctx, types.EventsOptions{Filters: filtersArgs})

	go func() {
		defer cancel()

		for {
			select {
			case <-ctx.Done():
				return

			case message := <-messages:
				messageJSON, err := json.Marshal(message)
				if err != nil {
					log.Printf(`[%s] Error while marshalling event: %v`, user.Username, err)
					return
				}

				if err = ws.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
					log.Printf(`[%s] Error while writing to events socket: %v`, user.Username, err)
					return
				}

			case err := <-errors:
				log.Printf(`[%s] Error while reading events: %v`, user.Username, err)
				return
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			break

		default:
			if readUntilClose(user, ws, `events`) {
				return
			}
		}
	}
}

func statsWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	stats, err := docker.ContainerStats(ctx, string(containerID), true)
	if err != nil {
		log.Print(err)
		return
	}
	defer stats.Body.Close()

	go func() {
		defer cancel()

		scanner := bufio.NewScanner(stats.Body)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				log.Printf(`[%s] Stats context is over for writing`, user.Username)
				return

			default:
				if err = ws.WriteMessage(websocket.TextMessage, scanner.Bytes()); err != nil {
					log.Printf(`[%s] Error while writing to stats socket: %v`, user.Username, err)
					return
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			log.Printf(`[%s] Stats context is over for reading`, user.Username)
			return

		default:
			if readUntilClose(user, ws, `stats`) {
				return
			}
		}
	}
}

func streamLogs(ctx context.Context, cancel context.CancelFunc, user *auth.User, containerID string, output chan<- []byte) {
	logs, err := docker.ContainerLogs(ctx, containerID, types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	defer cancel()

	if err != nil {
		log.Printf(`[%s] Logs opening in error: %v`, user.Username, err)
		return
	}
	defer logs.Close()

	scanner := bufio.NewScanner(logs)
	log.Printf(`[%s] Logs streaming started for %s`, user.Username, containerID)

	for scanner.Scan() {
		logLine := scanner.Bytes()
		if len(logLine) > ignoredByteLogSize {
			output <- logLine[ignoredByteLogSize:]
		}
	}

	log.Printf(`[%s] Logs streaming ended for %s`, user.Username, containerID)
}

func streamEvents(ctx context.Context, cancel context.CancelFunc, user *auth.User, _ string, output chan<- []byte) {
	defer cancel()

	filtersArgs := filters.NewArgs()
	if err := labelFilters(&filtersArgs, user, nil); err != nil {
		log.Printf(`[%s] Events opening in error: %v`, user.Username, err)
		return
	}
	if err := eventFilters(&filtersArgs); err != nil {
		log.Printf(`[%s] Events opening in error: %v`, user.Username, err)
		return
	}

	messages, errors := docker.Events(ctx, types.EventsOptions{Filters: filtersArgs})

	log.Printf(`[%s] Events streaming started`, user.Username)
	for {
		select {
		case <-ctx.Done():
			log.Printf(`[%s] Events streaming ended`, user.Username)
			return

		case message := <-messages:
			messageJSON, err := json.Marshal(message)
			if err != nil {
				log.Printf(`[%s] Events marshalling in error: %v`, user.Username, err)
				cancel()
			} else {
				output <- messageJSON
			}

		case err := <-errors:
			log.Printf(`[%s] Events reading in error: %v`, user.Username, err)
			cancel()
		}
	}
}

func streamStats(ctx context.Context, cancel context.CancelFunc, user *auth.User, containerID string, output chan<- []byte) {
	stats, err := docker.ContainerStats(ctx, containerID, true)
	defer cancel()

	if err != nil {
		log.Printf(`[%s] Stats opening in error for %s: %v`, user.Username, containerID, err)
		return
	}
	defer stats.Body.Close()

	scanner := bufio.NewScanner(stats.Body)
	log.Printf(`[%s] Stats streaming started for %s`, user.Username, containerID)

	for scanner.Scan() {
		output <- scanner.Bytes()
	}

	log.Printf(`[%s] Stats streaming ended for %s`, user.Username, containerID)
}

func handleBusDemand(user *auth.User, name string, input []byte, demand *regexp.Regexp, cancel context.CancelFunc, output chan<- []byte, streamFn func(context.Context, context.CancelFunc, *auth.User, string, chan<- []byte)) context.CancelFunc {
	demandGroups := demand.FindSubmatch(input)
	if len(demandGroups) < 2 {
		log.Printf(`[%s] Unable to parse bus demand %s for %s`, user.Username, input, name)
	}

	action := string(demandGroups[1])

	containerID := ``
	if len(demandGroups) > 2 {
		containerID = string(demandGroups[2])
	}

	if action == stop && cancel != nil {
		log.Printf(`[%s] Stopping %s stream`, user.Username, name)
		cancel()
	} else if action == start {
		log.Printf(`[%s] Starting %s stream for %s`, user.Username, name, containerID)

		if cancel != nil {
			log.Printf(`[%s] Cancelling previous %s stream`, user.Username, name)
			cancel()
		}

		ctx, newCancel := context.WithCancel(context.Background())
		go streamFn(ctx, newCancel, user, string(containerID), output)

		return newCancel
	}

	return nil
}

func busWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	done := make(chan int)

	output := make(chan []byte)
	defer close(output)

	input := make(chan []byte)
	defer close(input)

	go readContent(user, ws, `streaming`, done, input)
	log.Printf(`[%s] Streaming started`, user.Username)

	var eventsCancelFunc context.CancelFunc
	var logsCancelFunc context.CancelFunc
	var statsCancelFunc context.CancelFunc

	for {
		select {
		case <-done:
			log.Printf(`[%s] Streaming ended`, user.Username)
			return

		case inputBytes := <-input:
			if eventsDemand.Match(inputBytes) {
				eventsCancelFunc = handleBusDemand(user, `events`, inputBytes, eventsDemand, eventsCancelFunc, output, streamEvents)
				if eventsCancelFunc != nil {
					defer eventsCancelFunc()
				}
			} else if logsDemand.Match(inputBytes) {
				logsCancelFunc = handleBusDemand(user, `logs`, inputBytes, logsDemand, logsCancelFunc, output, streamLogs)
				if logsCancelFunc != nil {
					defer logsCancelFunc()
				}
			} else if statsDemand.Match(inputBytes) {
				statsCancelFunc = handleBusDemand(user, `stats`, inputBytes, statsDemand, statsCancelFunc, output, streamStats)
				if statsCancelFunc != nil {
					defer statsCancelFunc()
				}
			}

		case outputBytes := <-output:
			if err = ws.WriteMessage(websocket.TextMessage, outputBytes); err != nil {
				log.Printf(`[%s] Error while writing to streaming: %v`, user.Username, err)
				close(done)
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	urlPath := []byte(r.URL.Path)

	if logWebsocketRequest.Match(urlPath) {
		logsContainerWebsocketHandler(w, r, logWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if eventsWebsocketRequest.Match((urlPath)) {
		eventsWebsocketHandler(w, r)
	} else if statsWebsocketRequest.Match(urlPath) {
		statsWebsocketHandler(w, r, statsWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if busWebsocketRequest.Match(urlPath) {
		busWebsocketHandler(w, r)
	}
}
<|endoftext|>


<commit_before>package spec

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/mackerelio/golib/logging"
)

// This Generator collects metadata about cloud instances.
// Currently EC2 and GCE are supported.
// EC2: http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/AESDG-chapter-instancedata.html
// GCE: https://developers.google.com/compute/docs/metadata

// CloudGenerator definition
type CloudGenerator struct {
	CloudMetaGenerator
}

// CloudMetaGenerator interface of metadata generator for each cloud platform
type CloudMetaGenerator interface {
	Generate() (interface{}, error)
	SuggestCustomIdentifier() (string, error)
}

// Key is a root key for the generator.
func (g *CloudGenerator) Key() string {
	return "cloud"
}

var cloudLogger = logging.GetLogger("spec.cloud")

var ec2BaseURL, gceMetaURL *url.URL

func init() {
	ec2BaseURL, _ = url.Parse("http://169.254.169.254/latest/meta-data")
	gceMetaURL, _ = url.Parse("http://metadata.google.internal/computeMetadata/v1/?recursive=true")
}

var timeout = 100 * time.Millisecond

// SuggestCloudGenerator returns suitable CloudGenerator
func SuggestCloudGenerator() *CloudGenerator {
	if isEC2() {
		return &CloudGenerator{&EC2Generator{ec2BaseURL}}
	}
	if isGCE() {
		return &CloudGenerator{&GCEGenerator{gceMetaURL}}
	}

	return nil
}

func httpCli() *http.Client {
	return &http.Client{
		Timeout: timeout,
		Transport: &http.Transport{
			// don't use HTTP_PROXY when requesting cloud instance metadata APIs
			Proxy: nil,
		},
	}
}

func isEC2() bool {
	cl := httpCli()
	// '/ami-id` is may be aws specific URL
	resp, err := cl.Get(ec2BaseURL.String() + "/ami-id")
	if err != nil {
		return false
	}
	defer resp.Body.Close()

	return resp.StatusCode == 200
}

func isGCE() bool {
	_, err := requestGCEMeta()
	return err == nil
}

func requestGCEMeta() ([]byte, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", gceMetaURL.String(), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Metadata-Flavor", "Google")

	resp, err := cl.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("failed to request gce meta. response code: %d", resp.StatusCode)
	}
	return ioutil.ReadAll(resp.Body)
}

// EC2Generator meta generator for EC2
type EC2Generator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *EC2Generator) Generate() (interface{}, error) {
	cl := httpCli()

	metadataKeys := []string{
		"instance-id",
		"instance-type",
		"placement/availability-zone",
		"security-groups",
		"ami-id",
		"hostname",
		"local-hostname",
		"public-hostname",
		"local-ipv4",
		"public-ipv4",
		"reservation-id",
	}

	metadata := make(map[string]string)

	for _, key := range metadataKeys {
		resp, err := cl.Get(g.baseURL.String() + "/" + key)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on EC2. Error while reading '%s'", key)
			return nil, nil
		}
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadata[key] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}

	results := make(map[string]interface{})
	results["provider"] = "ec2"
	results["metadata"] = metadata

	return results, nil
}

// SuggestCustomIdentifier suggests the identifier of the EC2 instance
func (g *EC2Generator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	key := "instance-id"
	resp, err := cl.Get(g.baseURL.String() + "/" + key)
	if err != nil {
		return "", fmt.Errorf("error while retrieving instance-id")
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request instance-id. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting instance-id cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".ec2.amazonaws.com", nil
}

// GCEGenerator generate for GCE
type GCEGenerator struct {
	metaURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *GCEGenerator) Generate() (interface{}, error) {
	bytes, err := requestGCEMeta()
	if err != nil {
		return nil, err
	}
	var data gceMeta
	json.Unmarshal(bytes, &data)
	return data.toGeneratorResults(), nil
}

type gceInstance struct {
	Zone         string
	InstanceType string `json:"machineType"`
	Hostname     string
	InstanceID   uint64 `json:"id"`
}

type gceProject struct {
	ProjectID        string
	NumericProjectID uint64
}

type gceMeta struct {
	Instance *gceInstance
	Project  *gceProject
}

func (g gceMeta) toGeneratorMeta() map[string]string {
	meta := make(map[string]string)

	lastS := func(s string) string {
		ss := strings.Split(s, "/")
		return ss[len(ss)-1]
	}

	if ins := g.Instance; ins != nil {
		meta["hostname"] = ins.Hostname
		meta["instance-id"] = fmt.Sprint(ins.InstanceID)
		meta["instance-type"] = lastS(ins.InstanceType)
		meta["zone"] = lastS(ins.Zone)
	}

	if proj := g.Project; proj != nil {
		meta["projectId"] = proj.ProjectID
	}

	return meta
}

func (g gceMeta) toGeneratorResults() interface{} {
	results := make(map[string]interface{})
	results["provider"] = "gce"
	results["metadata"] = g.toGeneratorMeta()

	return results
}

// SuggestCustomIdentifier for GCE is not implemented yet
func (g *GCEGenerator) SuggestCustomIdentifier() (string, error) {
	return "", nil
}
<commit_msg>start implementing azure vm metadata retrieval<commit_after>package spec

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/mackerelio/golib/logging"
)

// This Generator collects metadata about cloud instances.
// Currently EC2 and GCE are supported.
// EC2: http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/AESDG-chapter-instancedata.html
// GCE: https://developers.google.com/compute/docs/metadata

// CloudGenerator definition
type CloudGenerator struct {
	CloudMetaGenerator
}

// CloudMetaGenerator interface of metadata generator for each cloud platform
type CloudMetaGenerator interface {
	Generate() (interface{}, error)
	SuggestCustomIdentifier() (string, error)
}

// Key is a root key for the generator.
func (g *CloudGenerator) Key() string {
	return "cloud"
}

var cloudLogger = logging.GetLogger("spec.cloud")

var ec2BaseURL, gceMetaURL, azureVMBaseURL *url.URL

func init() {
	ec2BaseURL, _ = url.Parse("http://169.254.169.254/latest/meta-data")
	gceMetaURL, _ = url.Parse("http://metadata.google.internal/computeMetadata/v1/?recursive=true")
	azureVMBaseURL, _ = url.Parse("http://169.254.169.254/metadata/instance")
}

var timeout = 100 * time.Millisecond

// SuggestCloudGenerator returns suitable CloudGenerator
func SuggestCloudGenerator() *CloudGenerator {
	if isEC2() {
		return &CloudGenerator{&EC2Generator{ec2BaseURL}}
	}
	if isGCE() {
		return &CloudGenerator{&GCEGenerator{gceMetaURL}}
	}

	return nil
}

func httpCli() *http.Client {
	return &http.Client{
		Timeout: timeout,
		Transport: &http.Transport{
			// don't use HTTP_PROXY when requesting cloud instance metadata APIs
			Proxy: nil,
		},
	}
}

func isEC2() bool {
	cl := httpCli()
	// '/ami-id` is may be aws specific URL
	resp, err := cl.Get(ec2BaseURL.String() + "/ami-id")
	if err != nil {
		return false
	}
	defer resp.Body.Close()

	return resp.StatusCode == 200
}

func isGCE() bool {
	_, err := requestGCEMeta()
	return err == nil
}

func isAzure() bool {
	cl := httpCli()
	// '/vmId` is probably Azure VM specific URL
	resp, err := cl.Get(azureVMBaseURL.String() + "/compute/vmId?api-version=2017-04-02&format=text")
	if err != nil {
		return false
	}
	defer resp.Body.Close()

	return resp.StatusCode == 200
}

func requestGCEMeta() ([]byte, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", gceMetaURL.String(), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Metadata-Flavor", "Google")

	resp, err := cl.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("failed to request gce meta. response code: %d", resp.StatusCode)
	}
	return ioutil.ReadAll(resp.Body)
}

// EC2Generator meta generator for EC2
type EC2Generator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *EC2Generator) Generate() (interface{}, error) {
	cl := httpCli()

	metadataKeys := []string{
		"instance-id",
		"instance-type",
		"placement/availability-zone",
		"security-groups",
		"ami-id",
		"hostname",
		"local-hostname",
		"public-hostname",
		"local-ipv4",
		"public-ipv4",
		"reservation-id",
	}

	metadata := make(map[string]string)

	for _, key := range metadataKeys {
		resp, err := cl.Get(g.baseURL.String() + "/" + key)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on EC2. Error while reading '%s'", key)
			return nil, nil
		}
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadata[key] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}

	results := make(map[string]interface{})
	results["provider"] = "ec2"
	results["metadata"] = metadata

	return results, nil
}

// SuggestCustomIdentifier suggests the identifier of the EC2 instance
func (g *EC2Generator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	key := "instance-id"
	resp, err := cl.Get(g.baseURL.String() + "/" + key)
	if err != nil {
		return "", fmt.Errorf("error while retrieving instance-id")
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request instance-id. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting instance-id cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".ec2.amazonaws.com", nil
}

// GCEGenerator generate for GCE
type GCEGenerator struct {
	metaURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *GCEGenerator) Generate() (interface{}, error) {
	bytes, err := requestGCEMeta()
	if err != nil {
		return nil, err
	}
	var data gceMeta
	json.Unmarshal(bytes, &data)
	return data.toGeneratorResults(), nil
}

type gceInstance struct {
	Zone         string
	InstanceType string `json:"machineType"`
	Hostname     string
	InstanceID   uint64 `json:"id"`
}

type gceProject struct {
	ProjectID        string
	NumericProjectID uint64
}

type gceMeta struct {
	Instance *gceInstance
	Project  *gceProject
}

func (g gceMeta) toGeneratorMeta() map[string]string {
	meta := make(map[string]string)

	lastS := func(s string) string {
		ss := strings.Split(s, "/")
		return ss[len(ss)-1]
	}

	if ins := g.Instance; ins != nil {
		meta["hostname"] = ins.Hostname
		meta["instance-id"] = fmt.Sprint(ins.InstanceID)
		meta["instance-type"] = lastS(ins.InstanceType)
		meta["zone"] = lastS(ins.Zone)
	}

	if proj := g.Project; proj != nil {
		meta["projectId"] = proj.ProjectID
	}

	return meta
}

func (g gceMeta) toGeneratorResults() interface{} {
	results := make(map[string]interface{})
	results["provider"] = "gce"
	results["metadata"] = g.toGeneratorMeta()

	return results
}

// SuggestCustomIdentifier for GCE is not implemented yet
func (g *GCEGenerator) SuggestCustomIdentifier() (string, error) {
	return "", nil
}

// AzureVMGenerator meta generator for Azure VM
type AzureVMGenerator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *AzureVMGenerator) Generate() (interface{}, error) {
	metadataComputeKeys := []string{
		"location",
		"name",
		"offer",
		"osType",
		"platformFaultDomain",
		"platformUpdateDomain",
		"publisher",
		"sku",
		"version",
		"vmId",
		"vmSize",
	}

	ipAddressKeys := []string{
		"privateIpAddress",
		"publicIpAddress",
	}

	metadata := make(map[string]string)
	metadata = retrieveAzureVMMetadata(metadata, g.baseURL.String(), "/compute/", metadataComputeKeys)
	metadata = retrieveAzureVMMetadata(metadata, g.baseURL.String(), "/network/interface/0/ipAddress/", ipAddressKeys)

	results := make(map[string]interface{})
	results["provider"] = "AzureVM"
	results["metadata"] = metadata

	return results, nil
}

func retrieveAzureVMMetadata(metadataMap map[string]string, baseURL string, urlSuffix string, keys []string) map[string]string {
	cl := httpCli()

	for _, key := range keys {
		resp, err := cl.Get(baseURL + urlSuffix + key + "?api-version=2017-04-02&format=text")
		if err != nil {
			cloudLogger.Debugf("This host may not be running on Azure VM. Error while reading '%s'", key)
			return nil
		}
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadataMap[key] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}
	return metadataMap
}

// SuggestCustomIdentifier suggests the identifier of the Azure VM instance
func (g *AzureVMGenerator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	resp, err := cl.Get(azureVMBaseURL.String() + "/compute/vmId?api-version=2017-04-02&format=text")
	if err != nil {
		return "", fmt.Errorf("error while retrieving vmId")
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request vmId. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting vmId cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return "virtual_machine." + instanceID + ".azure.microsoft.com", nil
}
<|endoftext|>


<commit_before>package admin

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/jinzhu/now"
	"github.com/qor/inflection"
	"github.com/qor/qor"
	"github.com/qor/qor/resource"
	"github.com/qor/qor/roles"
	"github.com/qor/qor/utils"
)

type Resource struct {
	resource.Resource
	admin         *Admin
	Config        *Config
	Metas         []*Meta
	actions       []*Action
	scopes        []*Scope
	filters       map[string]*Filter
	searchAttrs   *[]string
	sortableAttrs *[]string
	indexAttrs    []string
	newAttrs      []string
	editAttrs     []string
	showAttrs     []string
	cachedMetas   *map[string][]*Meta
	SearchHandler func(keyword string, context *qor.Context) *gorm.DB
}

func (res *Resource) Meta(meta *Meta) {
	if res.GetMeta(meta.Name) != nil {
		utils.ExitWithMsg("Duplicated meta %v defined for resource %v", meta.Name, res.Name)
	}

	meta.base = res
	meta.updateMeta()
	res.Metas = append(res.Metas, meta)
}

func (res Resource) GetAdmin() *Admin {
	return res.admin
}

func (res Resource) ToParam() string {
	if res.Config.Singleton == true {
		return utils.ToParamString(res.Name)
	}
	return utils.ToParamString(inflection.Plural(res.Name))
}

func (res Resource) UseTheme(theme string) []string {
	if res.Config != nil {
		res.Config.Themes = append(res.Config.Themes, theme)
		return res.Config.Themes
	}
	return []string{}
}

func (res *Resource) convertObjectToMap(context *Context, value interface{}, kind string) interface{} {
	reflectValue := reflect.Indirect(reflect.ValueOf(value))
	switch reflectValue.Kind() {
	case reflect.Slice:
		values := []interface{}{}
		for i := 0; i < reflectValue.Len(); i++ {
			values = append(values, res.convertObjectToMap(context, reflectValue.Index(i).Interface(), kind))
		}
		return values
	case reflect.Struct:
		var metas []*Meta
		if kind == "index" {
			metas = res.indexMetas()
		} else if kind == "show" {
			metas = res.showMetas()
		}

		values := map[string]interface{}{}
		for _, meta := range metas {
			if meta.HasPermission(roles.Read, context.Context) {
				if valuer := meta.GetValuer(); valuer != nil {
					value := valuer(value, context.Context)
					if meta.Resource != nil {
						value = meta.Resource.(*Resource).convertObjectToMap(context, value, kind)
					}
					values[meta.GetName()] = value
				}
			}
		}
		return values
	default:
		panic(fmt.Sprintf("Can't convert %v (%v) to map", reflectValue, reflectValue.Kind()))
	}
}

func (res *Resource) Decode(context *qor.Context, value interface{}) error {
	return resource.Decode(context, value, res)
}

func (res *Resource) allAttrs() []string {
	var attrs []string
	scope := &gorm.Scope{Value: res.Value}

Fields:
	for _, field := range scope.GetModelStruct().StructFields {
		for _, meta := range res.Metas {
			if field.Name == meta.FieldName {
				attrs = append(attrs, meta.Name)
				continue Fields
			}
		}

		if field.IsForeignKey {
			continue
		}

		for _, value := range []string{"CreatedAt", "UpdatedAt", "DeletedAt"} {
			if value == field.Name {
				continue Fields
			}
		}

		if (field.IsNormal || field.Relationship != nil) && !field.IsIgnored {
			attrs = append(attrs, field.Name)
		}
	}

MetaIncluded:
	for _, meta := range res.Metas {
		for _, attr := range attrs {
			if attr == meta.FieldName || attr == meta.Name {
				continue MetaIncluded
			}
		}
		attrs = append(attrs, meta.Name)
	}

	return attrs
}

func (res *Resource) getAttrs(attrs []string) []string {
	if len(attrs) == 0 {
		return res.allAttrs()
	} else {
		var onlyExcludeAttrs = true
		for _, attr := range attrs {
			if !strings.HasPrefix(attr, "-") {
				onlyExcludeAttrs = false
				break
			}
		}
		if onlyExcludeAttrs {
			return append(res.allAttrs(), attrs...)
		}
		return attrs
	}
}

func (res *Resource) IndexAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.indexAttrs = columns
	}
	return res.getAttrs(res.indexAttrs)
}

func (res *Resource) NewAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.newAttrs = columns
	}
	return res.getAttrs(res.newAttrs)
}

func (res *Resource) EditAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.editAttrs = columns
	}
	return res.getAttrs(res.editAttrs)
}

func (res *Resource) ShowAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.showAttrs = columns
	}
	return res.getAttrs(res.showAttrs)
}

func (res *Resource) SortableAttrs(columns ...string) []string {
	if len(columns) != 0 || res.sortableAttrs == nil {
		if len(columns) == 0 {
			columns = res.indexAttrs
		}
		res.sortableAttrs = &[]string{}
		scope := res.GetAdmin().Config.DB.NewScope(res.Value)
		for _, column := range columns {
			if field, ok := scope.FieldByName(column); ok && field.DBName != "" {
				attrs := append(*res.sortableAttrs, column)
				res.sortableAttrs = &attrs
			}
		}
	}
	return *res.sortableAttrs
}

func (res *Resource) SearchAttrs(columns ...string) []string {
	if len(columns) != 0 || res.searchAttrs == nil {
		if len(columns) == 0 {
			columns = res.IndexAttrs()
		}

		res.searchAttrs = &columns
		res.SearchHandler = func(keyword string, context *qor.Context) *gorm.DB {
			db := context.GetDB()
			var joinConditionsMap = map[string][]string{}
			var conditions []string
			var keywords []interface{}
			scope := db.NewScope(res.Value)

			for _, column := range columns {
				currentScope, nextScope := scope, scope

				if strings.Contains(column, ".") {
					for _, field := range strings.Split(column, ".") {
						column = field
						currentScope = nextScope
						if field, ok := scope.FieldByName(field); ok {
							if relationship := field.Relationship; relationship != nil {
								nextScope = currentScope.New(reflect.New(field.Field.Type()).Interface())
								key := fmt.Sprintf("LEFT JOIN %v ON", nextScope.TableName())

								for index := range relationship.ForeignDBNames {
									if relationship.Kind == "has_one" || relationship.Kind == "has_many" {
										joinConditionsMap[key] = append(joinConditionsMap[key],
											fmt.Sprintf("%v.%v = %v.%v",
												nextScope.QuotedTableName(), scope.Quote(relationship.ForeignDBNames[index]),
												currentScope.QuotedTableName(), scope.Quote(relationship.AssociationForeignDBNames[index]),
											))
									} else if relationship.Kind == "belongs_to" {
										joinConditionsMap[key] = append(joinConditionsMap[key],
											fmt.Sprintf("%v.%v = %v.%v",
												currentScope.QuotedTableName(), scope.Quote(relationship.ForeignDBNames[index]),
												nextScope.QuotedTableName(), scope.Quote(relationship.AssociationForeignDBNames[index]),
											))
									}
								}
							}
						}
					}
				}

				var tableName = currentScope.Quote(currentScope.TableName())
				if field, ok := currentScope.FieldByName(column); ok && field.IsNormal {
					switch field.Field.Kind() {
					case reflect.String:
						conditions = append(conditions, fmt.Sprintf("upper(%v.%v) like upper(?)", tableName, scope.Quote(field.DBName)))
						keywords = append(keywords, "%"+keyword+"%")
					case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
						if _, err := strconv.Atoi(keyword); err == nil {
							conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
							keywords = append(keywords, keyword)
						}
					case reflect.Float32, reflect.Float64:
						if _, err := strconv.ParseFloat(keyword, 64); err == nil {
							conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
							keywords = append(keywords, keyword)
						}
					case reflect.Bool:
						if value, err := strconv.ParseBool(keyword); err == nil {
							conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
							keywords = append(keywords, value)
						}
					case reflect.Struct:
						// time ?
						if _, ok := field.Field.Interface().(time.Time); ok {
							if parsedTime, err := now.Parse(keyword); err == nil {
								conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
								keywords = append(keywords, parsedTime)
							}
						}
					case reflect.Ptr:
						// time ?
						if _, ok := field.Field.Interface().(*time.Time); ok {
							if parsedTime, err := now.Parse(keyword); err == nil {
								conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
								keywords = append(keywords, parsedTime)
							}
						}
					default:
						conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
						keywords = append(keywords, keyword)
					}
				}
			}

			// join conditions
			if len(joinConditionsMap) > 0 {
				var joinConditions []string
				for key, values := range joinConditionsMap {
					joinConditions = append(joinConditions, fmt.Sprintf("%v %v", key, strings.Join(values, " AND ")))
				}
				db = db.Joins(strings.Join(joinConditions, " "))
			}

			// search conditions
			if len(conditions) > 0 {
				return db.Where(strings.Join(conditions, " OR "), keywords...)
			} else {
				return db
			}
		}
	}

	return *res.searchAttrs
}

func (res *Resource) getCachedMetas(cacheKey string, fc func() []resource.Metaor) []*Meta {
	if res.cachedMetas == nil {
		res.cachedMetas = &map[string][]*Meta{}
	}

	if values, ok := (*res.cachedMetas)[cacheKey]; ok {
		return values
	} else {
		values := fc()
		var metas []*Meta
		for _, value := range values {
			metas = append(metas, value.(*Meta))
		}
		(*res.cachedMetas)[cacheKey] = metas
		return metas
	}
}

func (res *Resource) GetMetas(attrs []string) []resource.Metaor {
	if len(attrs) == 0 {
		attrs = res.allAttrs()
	}
	var showAttrs, ignoredAttrs []string
	for _, attr := range attrs {
		if strings.HasPrefix(attr, "-") {
			ignoredAttrs = append(ignoredAttrs, strings.TrimLeft(attr, "-"))
		} else {
			showAttrs = append(showAttrs, attr)
		}
	}

	primaryKey := res.PrimaryFieldName()

	metas := []resource.Metaor{}
Attrs:
	for _, attr := range showAttrs {
		for _, a := range ignoredAttrs {
			if attr == a {
				continue Attrs
			}
		}

		var meta *Meta
		for _, m := range res.Metas {
			if m.GetName() == attr {
				meta = m
				break
			}
		}

		if meta == nil {
			meta = &Meta{}
			meta.Name = attr
			meta.base = res
			if attr == primaryKey {
				meta.Type = "hidden"
			}
			meta.updateMeta()
		}

		metas = append(metas, meta)
	}

	return metas
}

func (res *Resource) GetMeta(name string) *Meta {
	for _, meta := range res.Metas {
		if meta.Name == name || meta.GetFieldName() == name {
			return meta
		}
	}
	return nil
}

func (res *Resource) indexMetas() []*Meta {
	return res.getCachedMetas("index_metas", func() []resource.Metaor {
		return res.GetMetas(res.IndexAttrs())
	})
}

func (res *Resource) newMetas() []*Meta {
	return res.getCachedMetas("new_metas", func() []resource.Metaor {
		return res.GetMetas(res.NewAttrs())
	})
}

func (res *Resource) editMetas() []*Meta {
	return res.getCachedMetas("edit_metas", func() []resource.Metaor {
		return res.GetMetas(res.EditAttrs())
	})
}

func (res *Resource) showMetas() []*Meta {
	return res.getCachedMetas("show_metas", func() []resource.Metaor {
		return res.GetMetas(res.ShowAttrs())
	})
}

func (res *Resource) allMetas() []*Meta {
	return res.getCachedMetas("all_metas", func() []resource.Metaor {
		return res.GetMetas([]string{})
	})
}

func (res *Resource) allowedMetas(attrs []*Meta, context *Context, roles ...roles.PermissionMode) []*Meta {
	var metas = []*Meta{}
	for _, meta := range attrs {
		for _, role := range roles {
			if meta.HasPermission(role, context.Context) {
				metas = append(metas, meta)
				break
			}
		}
	}
	return metas
}
<commit_msg>Only generate search handler if has valid search attrs<commit_after>package admin

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/jinzhu/now"
	"github.com/qor/inflection"
	"github.com/qor/qor"
	"github.com/qor/qor/resource"
	"github.com/qor/qor/roles"
	"github.com/qor/qor/utils"
)

type Resource struct {
	resource.Resource
	admin         *Admin
	Config        *Config
	Metas         []*Meta
	actions       []*Action
	scopes        []*Scope
	filters       map[string]*Filter
	searchAttrs   *[]string
	sortableAttrs *[]string
	indexAttrs    []string
	newAttrs      []string
	editAttrs     []string
	showAttrs     []string
	cachedMetas   *map[string][]*Meta
	SearchHandler func(keyword string, context *qor.Context) *gorm.DB
}

func (res *Resource) Meta(meta *Meta) {
	if res.GetMeta(meta.Name) != nil {
		utils.ExitWithMsg("Duplicated meta %v defined for resource %v", meta.Name, res.Name)
	}

	meta.base = res
	meta.updateMeta()
	res.Metas = append(res.Metas, meta)
}

func (res Resource) GetAdmin() *Admin {
	return res.admin
}

func (res Resource) ToParam() string {
	if res.Config.Singleton == true {
		return utils.ToParamString(res.Name)
	}
	return utils.ToParamString(inflection.Plural(res.Name))
}

func (res Resource) UseTheme(theme string) []string {
	if res.Config != nil {
		res.Config.Themes = append(res.Config.Themes, theme)
		return res.Config.Themes
	}
	return []string{}
}

func (res *Resource) convertObjectToMap(context *Context, value interface{}, kind string) interface{} {
	reflectValue := reflect.Indirect(reflect.ValueOf(value))
	switch reflectValue.Kind() {
	case reflect.Slice:
		values := []interface{}{}
		for i := 0; i < reflectValue.Len(); i++ {
			values = append(values, res.convertObjectToMap(context, reflectValue.Index(i).Interface(), kind))
		}
		return values
	case reflect.Struct:
		var metas []*Meta
		if kind == "index" {
			metas = res.indexMetas()
		} else if kind == "show" {
			metas = res.showMetas()
		}

		values := map[string]interface{}{}
		for _, meta := range metas {
			if meta.HasPermission(roles.Read, context.Context) {
				if valuer := meta.GetValuer(); valuer != nil {
					value := valuer(value, context.Context)
					if meta.Resource != nil {
						value = meta.Resource.(*Resource).convertObjectToMap(context, value, kind)
					}
					values[meta.GetName()] = value
				}
			}
		}
		return values
	default:
		panic(fmt.Sprintf("Can't convert %v (%v) to map", reflectValue, reflectValue.Kind()))
	}
}

func (res *Resource) Decode(context *qor.Context, value interface{}) error {
	return resource.Decode(context, value, res)
}

func (res *Resource) allAttrs() []string {
	var attrs []string
	scope := &gorm.Scope{Value: res.Value}

Fields:
	for _, field := range scope.GetModelStruct().StructFields {
		for _, meta := range res.Metas {
			if field.Name == meta.FieldName {
				attrs = append(attrs, meta.Name)
				continue Fields
			}
		}

		if field.IsForeignKey {
			continue
		}

		for _, value := range []string{"CreatedAt", "UpdatedAt", "DeletedAt"} {
			if value == field.Name {
				continue Fields
			}
		}

		if (field.IsNormal || field.Relationship != nil) && !field.IsIgnored {
			attrs = append(attrs, field.Name)
		}
	}

MetaIncluded:
	for _, meta := range res.Metas {
		for _, attr := range attrs {
			if attr == meta.FieldName || attr == meta.Name {
				continue MetaIncluded
			}
		}
		attrs = append(attrs, meta.Name)
	}

	return attrs
}

func (res *Resource) getAttrs(attrs []string) []string {
	if len(attrs) == 0 {
		return res.allAttrs()
	} else {
		var onlyExcludeAttrs = true
		for _, attr := range attrs {
			if !strings.HasPrefix(attr, "-") {
				onlyExcludeAttrs = false
				break
			}
		}
		if onlyExcludeAttrs {
			return append(res.allAttrs(), attrs...)
		}
		return attrs
	}
}

func (res *Resource) IndexAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.indexAttrs = columns
	}
	return res.getAttrs(res.indexAttrs)
}

func (res *Resource) NewAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.newAttrs = columns
	}
	return res.getAttrs(res.newAttrs)
}

func (res *Resource) EditAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.editAttrs = columns
	}
	return res.getAttrs(res.editAttrs)
}

func (res *Resource) ShowAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.showAttrs = columns
	}
	return res.getAttrs(res.showAttrs)
}

func (res *Resource) SortableAttrs(columns ...string) []string {
	if len(columns) != 0 || res.sortableAttrs == nil {
		if len(columns) == 0 {
			columns = res.indexAttrs
		}
		res.sortableAttrs = &[]string{}
		scope := res.GetAdmin().Config.DB.NewScope(res.Value)
		for _, column := range columns {
			if field, ok := scope.FieldByName(column); ok && field.DBName != "" {
				attrs := append(*res.sortableAttrs, column)
				res.sortableAttrs = &attrs
			}
		}
	}
	return *res.sortableAttrs
}

func (res *Resource) SearchAttrs(columns ...string) []string {
	if len(columns) != 0 || res.searchAttrs == nil {
		if len(columns) == 0 {
			columns = res.IndexAttrs()
		}

		if len(columns) > 0 {
			res.searchAttrs = &columns
			res.SearchHandler = func(keyword string, context *qor.Context) *gorm.DB {
				db := context.GetDB()
				var joinConditionsMap = map[string][]string{}
				var conditions []string
				var keywords []interface{}
				scope := db.NewScope(res.Value)

				for _, column := range columns {
					currentScope, nextScope := scope, scope

					if strings.Contains(column, ".") {
						for _, field := range strings.Split(column, ".") {
							column = field
							currentScope = nextScope
							if field, ok := scope.FieldByName(field); ok {
								if relationship := field.Relationship; relationship != nil {
									nextScope = currentScope.New(reflect.New(field.Field.Type()).Interface())
									key := fmt.Sprintf("LEFT JOIN %v ON", nextScope.TableName())

									for index := range relationship.ForeignDBNames {
										if relationship.Kind == "has_one" || relationship.Kind == "has_many" {
											joinConditionsMap[key] = append(joinConditionsMap[key],
												fmt.Sprintf("%v.%v = %v.%v",
													nextScope.QuotedTableName(), scope.Quote(relationship.ForeignDBNames[index]),
													currentScope.QuotedTableName(), scope.Quote(relationship.AssociationForeignDBNames[index]),
												))
										} else if relationship.Kind == "belongs_to" {
											joinConditionsMap[key] = append(joinConditionsMap[key],
												fmt.Sprintf("%v.%v = %v.%v",
													currentScope.QuotedTableName(), scope.Quote(relationship.ForeignDBNames[index]),
													nextScope.QuotedTableName(), scope.Quote(relationship.AssociationForeignDBNames[index]),
												))
										}
									}
								}
							}
						}
					}

					var tableName = currentScope.Quote(currentScope.TableName())
					if field, ok := currentScope.FieldByName(column); ok && field.IsNormal {
						switch field.Field.Kind() {
						case reflect.String:
							conditions = append(conditions, fmt.Sprintf("upper(%v.%v) like upper(?)", tableName, scope.Quote(field.DBName)))
							keywords = append(keywords, "%"+keyword+"%")
						case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
							if _, err := strconv.Atoi(keyword); err == nil {
								conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
								keywords = append(keywords, keyword)
							}
						case reflect.Float32, reflect.Float64:
							if _, err := strconv.ParseFloat(keyword, 64); err == nil {
								conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
								keywords = append(keywords, keyword)
							}
						case reflect.Bool:
							if value, err := strconv.ParseBool(keyword); err == nil {
								conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
								keywords = append(keywords, value)
							}
						case reflect.Struct:
							// time ?
							if _, ok := field.Field.Interface().(time.Time); ok {
								if parsedTime, err := now.Parse(keyword); err == nil {
									conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
									keywords = append(keywords, parsedTime)
								}
							}
						case reflect.Ptr:
							// time ?
							if _, ok := field.Field.Interface().(*time.Time); ok {
								if parsedTime, err := now.Parse(keyword); err == nil {
									conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
									keywords = append(keywords, parsedTime)
								}
							}
						default:
							conditions = append(conditions, fmt.Sprintf("%v.%v = ?", tableName, scope.Quote(field.DBName)))
							keywords = append(keywords, keyword)
						}
					}
				}

				// join conditions
				if len(joinConditionsMap) > 0 {
					var joinConditions []string
					for key, values := range joinConditionsMap {
						joinConditions = append(joinConditions, fmt.Sprintf("%v %v", key, strings.Join(values, " AND ")))
					}
					db = db.Joins(strings.Join(joinConditions, " "))
				}

				// search conditions
				if len(conditions) > 0 {
					return db.Where(strings.Join(conditions, " OR "), keywords...)
				} else {
					return db
				}
			}
		}
	}

	return columns
}

func (res *Resource) getCachedMetas(cacheKey string, fc func() []resource.Metaor) []*Meta {
	if res.cachedMetas == nil {
		res.cachedMetas = &map[string][]*Meta{}
	}

	if values, ok := (*res.cachedMetas)[cacheKey]; ok {
		return values
	} else {
		values := fc()
		var metas []*Meta
		for _, value := range values {
			metas = append(metas, value.(*Meta))
		}
		(*res.cachedMetas)[cacheKey] = metas
		return metas
	}
}

func (res *Resource) GetMetas(attrs []string) []resource.Metaor {
	if len(attrs) == 0 {
		attrs = res.allAttrs()
	}
	var showAttrs, ignoredAttrs []string
	for _, attr := range attrs {
		if strings.HasPrefix(attr, "-") {
			ignoredAttrs = append(ignoredAttrs, strings.TrimLeft(attr, "-"))
		} else {
			showAttrs = append(showAttrs, attr)
		}
	}

	primaryKey := res.PrimaryFieldName()

	metas := []resource.Metaor{}
Attrs:
	for _, attr := range showAttrs {
		for _, a := range ignoredAttrs {
			if attr == a {
				continue Attrs
			}
		}

		var meta *Meta
		for _, m := range res.Metas {
			if m.GetName() == attr {
				meta = m
				break
			}
		}

		if meta == nil {
			meta = &Meta{}
			meta.Name = attr
			meta.base = res
			if attr == primaryKey {
				meta.Type = "hidden"
			}
			meta.updateMeta()
		}

		metas = append(metas, meta)
	}

	return metas
}

func (res *Resource) GetMeta(name string) *Meta {
	for _, meta := range res.Metas {
		if meta.Name == name || meta.GetFieldName() == name {
			return meta
		}
	}
	return nil
}

func (res *Resource) indexMetas() []*Meta {
	return res.getCachedMetas("index_metas", func() []resource.Metaor {
		return res.GetMetas(res.IndexAttrs())
	})
}

func (res *Resource) newMetas() []*Meta {
	return res.getCachedMetas("new_metas", func() []resource.Metaor {
		return res.GetMetas(res.NewAttrs())
	})
}

func (res *Resource) editMetas() []*Meta {
	return res.getCachedMetas("edit_metas", func() []resource.Metaor {
		return res.GetMetas(res.EditAttrs())
	})
}

func (res *Resource) showMetas() []*Meta {
	return res.getCachedMetas("show_metas", func() []resource.Metaor {
		return res.GetMetas(res.ShowAttrs())
	})
}

func (res *Resource) allMetas() []*Meta {
	return res.getCachedMetas("all_metas", func() []resource.Metaor {
		return res.GetMetas([]string{})
	})
}

func (res *Resource) allowedMetas(attrs []*Meta, context *Context, roles ...roles.PermissionMode) []*Meta {
	var metas = []*Meta{}
	for _, meta := range attrs {
		for _, role := range roles {
			if meta.HasPermission(role, context.Context) {
				metas = append(metas, meta)
				break
			}
		}
	}
	return metas
}
<|endoftext|>


<commit_before>package k8sTest

import (
	"fmt"
	"strings"

	. "github.com/cilium/cilium/test/ginkgo-ext"
	"github.com/cilium/cilium/test/helpers"
	. "github.com/onsi/gomega"
)

var _ = Describe("K8sUpdates", func() {

	// This test runs 8 steps as following:
	// 1 - delete all pods. Clean cilium, this can be, and should be achieved by
	// `clean-cilium-state: "true"` option that we have in configmap
	// 2 - install cilium `cilium:v1.1.4`
	// 3 - make endpoints talk with each other with policy
	// 4 - upgrade cilium to `k8s1:5000/cilium/cilium-dev:latest`
	// 5 - make endpoints talk with each other with policy
	// 6 - downgrade cilium to `cilium:v1.1.4`
	// 7 - make endpoints talk with each other with policy
	// 8 - delete all pods. Clean cilium, this can be, and should be achieved by
	// `clean-cilium-state: "true"` option that we have in configmap.
	// This makes sure the upgrade tests won't affect any other test
	// 9 - re install cilium:latest image for remaining tests.

	var (
		kubectl *helpers.Kubectl

		cleanupCallback = func() { return }
	)

	BeforeAll(func() {
		kubectl = helpers.CreateKubectl(helpers.K8s1VMName(), logger)

		_ = kubectl.Delete(helpers.DNSDeployment())

		// Delete kube-dns because if not will be a restore the old endpoints
		// from master instead of create the new ones.
		_ = kubectl.DeleteResource(
			"deploy", fmt.Sprintf("-n %s kube-dns", helpers.KubeSystemNamespace))

		// Sometimes PolicyGen has a lot of pods running around without delete
		// it. Using this we are sure that we delete before this test start
		kubectl.Exec(fmt.Sprintf(
			"%s delete --all pods,svc,cnp -n %s", helpers.KubectlCmd, helpers.DefaultNamespace))

		kubectl.DeleteETCDOperator()

		ExpectAllPodsTerminated(kubectl)
	})

	AfterAll(func() {
		_ = kubectl.Apply(helpers.DNSDeployment())
	})

	AfterFailed(func() {
		kubectl.CiliumReport(helpers.KubeSystemNamespace, "cilium endpoint list")
	})

	JustAfterEach(func() {
		kubectl.ValidateNoErrorsInLogs(CurrentGinkgoTestDescription().Duration)
	})

	AfterEach(func() {
		cleanupCallback()
		ExpectAllPodsTerminated(kubectl)
	})

	It("Tests upgrade and downgrade from a Cilium stable image to master", func() {
		var assertUpgradeSuccessful func()
		assertUpgradeSuccessful, cleanupCallback =
			InstallAndValidateCiliumUpgrades(kubectl, helpers.CiliumStableVersion, helpers.CiliumDeveloperImage)
		assertUpgradeSuccessful()
	})
})

// InstallAndValidateCiliumUpgrades installs and tests if the oldVersion can be
// upgrade to the newVersion and if the newVersion can be downgraded to the
// oldVersion.  It returns two callbacks, the first one is the assertfunction
// that need to run, and the second one are the cleanup actions
func InstallAndValidateCiliumUpgrades(kubectl *helpers.Kubectl, oldVersion, newVersion string) (func(), func()) {
	canRun, err := helpers.CanRunK8sVersion(oldVersion, helpers.GetCurrentK8SEnv())
	ExpectWithOffset(1, err).To(BeNil(), "Unable to get k8s constraints for %s", oldVersion)
	if !canRun {
		Skip(fmt.Sprintf(
			"Cilium %q is not supported in K8s %q. Skipping upgrade/downgrade tests.",
			oldVersion, helpers.GetCurrentK8SEnv()))
		return func() {}, func() {}
	}
	switch helpers.GetCurrentIntegration() {
	case helpers.CIIntegrationFlannel:
		Skip(fmt.Sprintf(
			"Cilium %q and %q mode are not supported in K8s %q. Skipping upgrade/downgrade tests.",
			oldVersion, helpers.CIIntegrationFlannel, helpers.GetCurrentK8SEnv()))
		return func() {}, func() {}
	}

	demoPath := helpers.ManifestGet("demo.yaml")
	l7Policy := helpers.ManifestGet("l7-policy.yaml")
	apps := []string{helpers.App1, helpers.App2, helpers.App3}
	app1Service := "app1-service"

	cleanupCallback := func() {
		kubectl.Delete(l7Policy)
		kubectl.Delete(demoPath)

		// make sure that Kubedns is deleted correctly
		_ = kubectl.Delete(helpers.DNSDeployment())

		kubectl.DeleteETCDOperator()

		ExpectAllPodsTerminated(kubectl)

		// make sure we clean everything up before doing any other test
		err := kubectl.CiliumInstall(
			helpers.CiliumDefaultDSPatch,
			"cilium-cm-patch-clean-cilium-state.yaml",
		)
		ExpectWithOffset(1, err).To(BeNil(), "Cilium %q was not able to be deployed", newVersion)
		err = kubectl.WaitForCiliumInitContainerToFinish()
		ExpectWithOffset(1, err).To(BeNil(), "Cilium %q was not able to be clean up environment", newVersion)

		_ = kubectl.DeleteResource(
			"ds", fmt.Sprintf("-n %s cilium", helpers.KubeSystemNamespace))
	}

	testfunc := func() {
		// Making sure that we deleted the  cilium ds. No assert message
		// because maybe is not present
		_ = kubectl.DeleteResource("ds", fmt.Sprintf("-n %s cilium", helpers.KubeSystemNamespace))

		// Delete kube-dns because if not will be a restore the old endpoints
		// from master instead of create the new ones.
		_ = kubectl.Delete(helpers.DNSDeployment())

		ExpectAllPodsTerminated(kubectl)

		By("Installing a cleaning state of Cilium")
		err = kubectl.CiliumInstallVersion(
			helpers.CiliumDefaultDSPatch,
			"cilium-cm-patch-clean-cilium-state.yaml",
			oldVersion,
		)
		Expect(err).To(BeNil(), "Cilium %q was not able to be deployed", oldVersion)

		By("Installing kube-dns")
		_ = kubectl.Apply(helpers.DNSDeployment())

		// Deploy the etcd operator
		By("Deploying etcd-operator")
		err = kubectl.DeployETCDOperator()
		Expect(err).To(BeNil(), "Unable to deploy etcd operator")

		// Cilium is only ready if kvstore is ready, the kvstore is ready if
		// kube-dns is running.
		By("Cilium %q is installed and running", oldVersion)
		ExpectCiliumReady(kubectl)

		ExpectETCDOperatorReady(kubectl)

		By("Installing Microscope")
		microscopeErr, microscopeCancel := kubectl.MicroscopeStart()
		ExpectWithOffset(1, microscopeErr).To(BeNil(), "Microscope cannot be started")
		defer microscopeCancel()

		validatedImage := func(image string) {
			By("Checking that installed image is %q", image)

			filter := `{.items[*].status.containerStatuses[0].image}`
			data, err := kubectl.GetPods(
				helpers.KubeSystemNamespace, "-l k8s-app=cilium").Filter(filter)
			ExpectWithOffset(1, err).To(BeNil(), "Cannot get cilium pods")

			for _, val := range strings.Split(data.String(), " ") {
				ExpectWithOffset(1, val).To(ContainSubstring(image), "Cilium image didn't update correctly")
			}
		}

		validateEndpointsConnection := func() {
			By("Validate that endpoints are ready before making any connection")
			err := kubectl.CiliumEndpointWaitReady()
			ExpectWithOffset(1, err).To(BeNil(), "Endpoints are not ready after timeout")

			ExpectKubeDNSReady(kubectl)

			err = kubectl.WaitForKubeDNSEntry(app1Service, helpers.DefaultNamespace)
			ExpectWithOffset(1, err).To(BeNil(), "DNS entry is not ready after timeout")

			appPods := helpers.GetAppPods(apps, helpers.DefaultNamespace, kubectl, "id")

			err = kubectl.WaitForKubeDNSEntry(app1Service, helpers.DefaultNamespace)
			ExpectWithOffset(1, err).To(BeNil(), "DNS entry is not ready after timeout")

			By("Making L7 requests between endpoints")
			res := kubectl.ExecPodCmd(
				helpers.DefaultNamespace, appPods[helpers.App2],
				helpers.CurlFail("http://%s/public", app1Service))
			ExpectWithOffset(1, res).Should(helpers.CMDSuccess(), "Cannot curl app1-service")

			res = kubectl.ExecPodCmd(
				helpers.DefaultNamespace, appPods[helpers.App2],
				helpers.CurlFail("http://%s/private", app1Service))
			ExpectWithOffset(1, res).ShouldNot(helpers.CMDSuccess(), "Expect a 403 from app1-service")
		}

		By("Creating some endpoints and L7 policy")
		res := kubectl.Apply(demoPath)
		ExpectWithOffset(1, res).To(helpers.CMDSuccess(), "cannot apply dempo application")

		err := kubectl.WaitforPods(helpers.DefaultNamespace, "-l zgroup=testapp", timeout)
		Expect(err).Should(BeNil(), "Test pods are not ready after timeout")

		ExpectKubeDNSReady(kubectl)

		_, err = kubectl.CiliumPolicyAction(
			helpers.KubeSystemNamespace, l7Policy, helpers.KubectlApply, timeout)
		Expect(err).Should(BeNil(), "cannot import l7 policy: %v", l7Policy)

		validateEndpointsConnection()

		By("Updating cilium to master image")

		waitForUpdateImage := func(image string) func() bool {
			return func() bool {
				pods, err := kubectl.GetCiliumPods(helpers.KubeSystemNamespace)
				if err != nil {
					return false
				}

				filter := `{.items[*].status.containerStatuses[0].image}`
				data, err := kubectl.GetPods(
					helpers.KubeSystemNamespace, "-l k8s-app=cilium").Filter(filter)
				if err != nil {
					return false
				}
				number := strings.Count(data.String(), image)
				if number == len(pods) {
					return true
				}
				log.Infof("Only '%v' of '%v' cilium pods updated to the new image",
					number, len(pods))
				return false
			}
		}

		By("Install Cilium pre-flight check DaemonSet")
		err = kubectl.CiliumPreFlightInstall(helpers.CiliumDefaultPreFlightPatch)
		ExpectWithOffset(1, err).To(BeNil(), "Cilium pre-flight %q was not able to be deployed", newVersion)
		ExpectCiliumPreFlightInstallReady(kubectl)

		// Once they are installed we can remove it
		By("Removing Cilium pre-flight check DaemonSet")
		kubectl.Delete(helpers.GetK8sDescriptor(helpers.CiliumDefaultPreFlight))

		err = kubectl.CiliumInstall(helpers.CiliumDefaultDSPatch, helpers.CiliumConfigMapPatch)
		ExpectWithOffset(1, err).To(BeNil(), "Cilium %q was not able to be deployed", newVersion)

		err = helpers.WithTimeout(
			waitForUpdateImage(newVersion),
			"Cilium Pods are not updating correctly",
			&helpers.TimeoutConfig{Timeout: timeout})
		ExpectWithOffset(1, err).To(BeNil(), "Pods are not updating")

		err = kubectl.WaitforPods(
			helpers.KubeSystemNamespace, "-l k8s-app=cilium", timeout)
		ExpectWithOffset(1, err).Should(BeNil(), "Cilium is not ready after timeout")

		validatedImage(newVersion)

		validateEndpointsConnection()

		By("Downgrading cilium to %s image", oldVersion)

		err = kubectl.CiliumInstallVersion(
			helpers.CiliumDefaultDSPatch,
			helpers.CiliumConfigMapPatch,
			oldVersion,
		)
		ExpectWithOffset(1, err).To(BeNil(), "Cilium %q was not able to be deployed", oldVersion)

		err = helpers.WithTimeout(
			waitForUpdateImage(oldVersion),
			"Cilium Pods are not updating correctly",
			&helpers.TimeoutConfig{Timeout: timeout})
		ExpectWithOffset(1, err).To(BeNil(), "Pods are not updating")

		err = kubectl.WaitforPods(
			helpers.KubeSystemNamespace, "-l k8s-app=cilium", timeout)
		ExpectWithOffset(1, err).Should(BeNil(), "Cilium is not ready after timeout")

		validatedImage(oldVersion)

		validateEndpointsConnection()

	}
	return testfunc, cleanupCallback
}
<commit_msg>Test: Validate that cilium is ready after upgrade.<commit_after>package k8sTest

import (
	"fmt"
	"strings"

	. "github.com/cilium/cilium/test/ginkgo-ext"
	"github.com/cilium/cilium/test/helpers"
	. "github.com/onsi/gomega"
)

var _ = Describe("K8sUpdates", func() {

	// This test runs 8 steps as following:
	// 1 - delete all pods. Clean cilium, this can be, and should be achieved by
	// `clean-cilium-state: "true"` option that we have in configmap
	// 2 - install cilium `cilium:v1.1.4`
	// 3 - make endpoints talk with each other with policy
	// 4 - upgrade cilium to `k8s1:5000/cilium/cilium-dev:latest`
	// 5 - make endpoints talk with each other with policy
	// 6 - downgrade cilium to `cilium:v1.1.4`
	// 7 - make endpoints talk with each other with policy
	// 8 - delete all pods. Clean cilium, this can be, and should be achieved by
	// `clean-cilium-state: "true"` option that we have in configmap.
	// This makes sure the upgrade tests won't affect any other test
	// 9 - re install cilium:latest image for remaining tests.

	var (
		kubectl *helpers.Kubectl

		cleanupCallback = func() { return }
	)

	BeforeAll(func() {
		kubectl = helpers.CreateKubectl(helpers.K8s1VMName(), logger)

		_ = kubectl.Delete(helpers.DNSDeployment())

		// Delete kube-dns because if not will be a restore the old endpoints
		// from master instead of create the new ones.
		_ = kubectl.DeleteResource(
			"deploy", fmt.Sprintf("-n %s kube-dns", helpers.KubeSystemNamespace))

		// Sometimes PolicyGen has a lot of pods running around without delete
		// it. Using this we are sure that we delete before this test start
		kubectl.Exec(fmt.Sprintf(
			"%s delete --all pods,svc,cnp -n %s", helpers.KubectlCmd, helpers.DefaultNamespace))

		kubectl.DeleteETCDOperator()

		ExpectAllPodsTerminated(kubectl)
	})

	AfterAll(func() {
		_ = kubectl.Apply(helpers.DNSDeployment())
	})

	AfterFailed(func() {
		kubectl.CiliumReport(helpers.KubeSystemNamespace, "cilium endpoint list")
	})

	JustAfterEach(func() {
		kubectl.ValidateNoErrorsInLogs(CurrentGinkgoTestDescription().Duration)
	})

	AfterEach(func() {
		cleanupCallback()
		ExpectAllPodsTerminated(kubectl)
	})

	It("Tests upgrade and downgrade from a Cilium stable image to master", func() {
		var assertUpgradeSuccessful func()
		assertUpgradeSuccessful, cleanupCallback =
			InstallAndValidateCiliumUpgrades(kubectl, helpers.CiliumStableVersion, helpers.CiliumDeveloperImage)
		assertUpgradeSuccessful()
	})
})

// InstallAndValidateCiliumUpgrades installs and tests if the oldVersion can be
// upgrade to the newVersion and if the newVersion can be downgraded to the
// oldVersion.  It returns two callbacks, the first one is the assertfunction
// that need to run, and the second one are the cleanup actions
func InstallAndValidateCiliumUpgrades(kubectl *helpers.Kubectl, oldVersion, newVersion string) (func(), func()) {
	canRun, err := helpers.CanRunK8sVersion(oldVersion, helpers.GetCurrentK8SEnv())
	ExpectWithOffset(1, err).To(BeNil(), "Unable to get k8s constraints for %s", oldVersion)
	if !canRun {
		Skip(fmt.Sprintf(
			"Cilium %q is not supported in K8s %q. Skipping upgrade/downgrade tests.",
			oldVersion, helpers.GetCurrentK8SEnv()))
		return func() {}, func() {}
	}
	switch helpers.GetCurrentIntegration() {
	case helpers.CIIntegrationFlannel:
		Skip(fmt.Sprintf(
			"Cilium %q and %q mode are not supported in K8s %q. Skipping upgrade/downgrade tests.",
			oldVersion, helpers.CIIntegrationFlannel, helpers.GetCurrentK8SEnv()))
		return func() {}, func() {}
	}

	demoPath := helpers.ManifestGet("demo.yaml")
	l7Policy := helpers.ManifestGet("l7-policy.yaml")
	apps := []string{helpers.App1, helpers.App2, helpers.App3}
	app1Service := "app1-service"

	cleanupCallback := func() {
		kubectl.Delete(l7Policy)
		kubectl.Delete(demoPath)

		// make sure that Kubedns is deleted correctly
		_ = kubectl.Delete(helpers.DNSDeployment())

		kubectl.DeleteETCDOperator()

		ExpectAllPodsTerminated(kubectl)

		// make sure we clean everything up before doing any other test
		err := kubectl.CiliumInstall(
			helpers.CiliumDefaultDSPatch,
			"cilium-cm-patch-clean-cilium-state.yaml",
		)
		ExpectWithOffset(1, err).To(BeNil(), "Cilium %q was not able to be deployed", newVersion)
		err = kubectl.WaitForCiliumInitContainerToFinish()
		ExpectWithOffset(1, err).To(BeNil(), "Cilium %q was not able to be clean up environment", newVersion)

		_ = kubectl.DeleteResource(
			"ds", fmt.Sprintf("-n %s cilium", helpers.KubeSystemNamespace))
	}

	testfunc := func() {
		// Making sure that we deleted the  cilium ds. No assert message
		// because maybe is not present
		_ = kubectl.DeleteResource("ds", fmt.Sprintf("-n %s cilium", helpers.KubeSystemNamespace))

		// Delete kube-dns because if not will be a restore the old endpoints
		// from master instead of create the new ones.
		_ = kubectl.Delete(helpers.DNSDeployment())

		ExpectAllPodsTerminated(kubectl)

		By("Installing a cleaning state of Cilium")
		err = kubectl.CiliumInstallVersion(
			helpers.CiliumDefaultDSPatch,
			"cilium-cm-patch-clean-cilium-state.yaml",
			oldVersion,
		)
		Expect(err).To(BeNil(), "Cilium %q was not able to be deployed", oldVersion)

		By("Installing kube-dns")
		_ = kubectl.Apply(helpers.DNSDeployment())

		// Deploy the etcd operator
		By("Deploying etcd-operator")
		err = kubectl.DeployETCDOperator()
		Expect(err).To(BeNil(), "Unable to deploy etcd operator")

		// Cilium is only ready if kvstore is ready, the kvstore is ready if
		// kube-dns is running.
		By("Cilium %q is installed and running", oldVersion)
		ExpectCiliumReady(kubectl)

		ExpectETCDOperatorReady(kubectl)

		By("Installing Microscope")
		microscopeErr, microscopeCancel := kubectl.MicroscopeStart()
		ExpectWithOffset(1, microscopeErr).To(BeNil(), "Microscope cannot be started")
		defer microscopeCancel()

		validatedImage := func(image string) {
			By("Checking that installed image is %q", image)

			filter := `{.items[*].status.containerStatuses[0].image}`
			data, err := kubectl.GetPods(
				helpers.KubeSystemNamespace, "-l k8s-app=cilium").Filter(filter)
			ExpectWithOffset(1, err).To(BeNil(), "Cannot get cilium pods")

			for _, val := range strings.Split(data.String(), " ") {
				ExpectWithOffset(1, val).To(ContainSubstring(image), "Cilium image didn't update correctly")
			}
		}

		validateEndpointsConnection := func() {
			By("Validate that endpoints are ready before making any connection")
			err := kubectl.CiliumEndpointWaitReady()
			ExpectWithOffset(1, err).To(BeNil(), "Endpoints are not ready after timeout")

			ExpectKubeDNSReady(kubectl)

			err = kubectl.WaitForKubeDNSEntry(app1Service, helpers.DefaultNamespace)
			ExpectWithOffset(1, err).To(BeNil(), "DNS entry is not ready after timeout")

			appPods := helpers.GetAppPods(apps, helpers.DefaultNamespace, kubectl, "id")

			err = kubectl.WaitForKubeDNSEntry(app1Service, helpers.DefaultNamespace)
			ExpectWithOffset(1, err).To(BeNil(), "DNS entry is not ready after timeout")

			By("Making L7 requests between endpoints")
			res := kubectl.ExecPodCmd(
				helpers.DefaultNamespace, appPods[helpers.App2],
				helpers.CurlFail("http://%s/public", app1Service))
			ExpectWithOffset(1, res).Should(helpers.CMDSuccess(), "Cannot curl app1-service")

			res = kubectl.ExecPodCmd(
				helpers.DefaultNamespace, appPods[helpers.App2],
				helpers.CurlFail("http://%s/private", app1Service))
			ExpectWithOffset(1, res).ShouldNot(helpers.CMDSuccess(), "Expect a 403 from app1-service")
		}

		By("Creating some endpoints and L7 policy")
		res := kubectl.Apply(demoPath)
		ExpectWithOffset(1, res).To(helpers.CMDSuccess(), "cannot apply dempo application")

		err := kubectl.WaitforPods(helpers.DefaultNamespace, "-l zgroup=testapp", timeout)
		Expect(err).Should(BeNil(), "Test pods are not ready after timeout")

		ExpectKubeDNSReady(kubectl)

		_, err = kubectl.CiliumPolicyAction(
			helpers.KubeSystemNamespace, l7Policy, helpers.KubectlApply, timeout)
		Expect(err).Should(BeNil(), "cannot import l7 policy: %v", l7Policy)

		validateEndpointsConnection()

		By("Updating cilium to master image")

		waitForUpdateImage := func(image string) func() bool {
			return func() bool {
				pods, err := kubectl.GetCiliumPods(helpers.KubeSystemNamespace)
				if err != nil {
					return false
				}

				filter := `{.items[*].status.containerStatuses[0].image}`
				data, err := kubectl.GetPods(
					helpers.KubeSystemNamespace, "-l k8s-app=cilium").Filter(filter)
				if err != nil {
					return false
				}
				number := strings.Count(data.String(), image)
				if number == len(pods) {
					return true
				}
				log.Infof("Only '%v' of '%v' cilium pods updated to the new image",
					number, len(pods))
				return false
			}
		}

		By("Install Cilium pre-flight check DaemonSet")
		err = kubectl.CiliumPreFlightInstall(helpers.CiliumDefaultPreFlightPatch)
		ExpectWithOffset(1, err).To(BeNil(), "Cilium pre-flight %q was not able to be deployed", newVersion)
		ExpectCiliumPreFlightInstallReady(kubectl)

		// Once they are installed we can remove it
		By("Removing Cilium pre-flight check DaemonSet")
		kubectl.Delete(helpers.GetK8sDescriptor(helpers.CiliumDefaultPreFlight))

		err = kubectl.CiliumInstall(helpers.CiliumDefaultDSPatch, helpers.CiliumConfigMapPatch)
		ExpectWithOffset(1, err).To(BeNil(), "Cilium %q was not able to be deployed", newVersion)

		err = helpers.WithTimeout(
			waitForUpdateImage(newVersion),
			"Cilium Pods are not updating correctly",
			&helpers.TimeoutConfig{Timeout: timeout})
		ExpectWithOffset(1, err).To(BeNil(), "Pods are not updating")

		err = kubectl.WaitforPods(
			helpers.KubeSystemNamespace, "-l k8s-app=cilium", timeout)
		ExpectWithOffset(1, err).Should(BeNil(), "Cilium is not ready after timeout")

		validatedImage(newVersion)
		ExpectCiliumReady(kubectl)

		validateEndpointsConnection()

		By("Downgrading cilium to %s image", oldVersion)

		err = kubectl.CiliumInstallVersion(
			helpers.CiliumDefaultDSPatch,
			helpers.CiliumConfigMapPatch,
			oldVersion,
		)
		ExpectWithOffset(1, err).To(BeNil(), "Cilium %q was not able to be deployed", oldVersion)

		err = helpers.WithTimeout(
			waitForUpdateImage(oldVersion),
			"Cilium Pods are not updating correctly",
			&helpers.TimeoutConfig{Timeout: timeout})
		ExpectWithOffset(1, err).To(BeNil(), "Pods are not updating")

		err = kubectl.WaitforPods(
			helpers.KubeSystemNamespace, "-l k8s-app=cilium", timeout)
		ExpectWithOffset(1, err).Should(BeNil(), "Cilium is not ready after timeout")

		validatedImage(oldVersion)

		validateEndpointsConnection()

	}
	return testfunc, cleanupCallback
}
<|endoftext|>


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

package num

import (
	"math"

	"github.com/cpmech/gosl/chk"
	"github.com/cpmech/gosl/fun"
)

// The algorithms below are based on [1]
// REFERENCES:
// [1] Press WH, Teukolsky SA, Vetterling WT, Fnannery BP (2007) Numerical Recipes: The Art of
//     Scientific Computing. Third Edition. Cambridge University Press. 1235p.

// QuadElementary defines the interface for elementary quadrature algorithms with refinement.
type QuadElementary interface {
	Init(f fun.Ss, a, b, eps float64) // The constructor takes as inputs f, the function or functor to be integrated between limits a and b, also input.
	Integrate() (float64, error)      // Returns the integral for the specified input data
}

// Trap structure is used for the trapezoidal integration rule with refinement.
type ElementaryTrapz struct {
	n    int     // current level of refinement.
	a, b float64 // limits
	s    float64 // current value of the integral
	eps  float64 // precision
	f    fun.Ss  // the function
}

// Init initialises Trap structure
func (o *ElementaryTrapz) Init(f fun.Ss, a, b, eps float64) {
	o.n = 0
	o.f = f
	o.a = a
	o.b = b
	o.eps = eps
}

// Next returns the nth stage of refinement of the extended trapezoidal rule. On the first call (n=1),
// R b the routine returns the crudest estimate of a f .x/dx. Subsequent calls set n=2,3,... and
// improve the accuracy by adding 2 n-2 additional interior points.
func (o *ElementaryTrapz) Next() (res float64, err error) {
	var x, sum, del float64
	var it, j, tnm int
	o.n += 1
	var fa, fb, fx float64
	if o.n == 1 {
		fa, err = o.f(o.a)
		if err != nil {
			return
		}
		fb, err = o.f(o.b)
		if err != nil {
			return
		}
		o.s = 0.5 * (o.b - o.a) * (fa + fb)
		return o.s, nil
	} else {
		for it, j = 1, 1; j < o.n-1; j++ {
			it *= 2
		}
		tnm = it
		del = (o.b - o.a) / float64(tnm)

		// This is the spacing of the points to be added.
		x = o.a + 0.5*del

		for sum, j = 0.0, 0; j < it; j, x = j+1, x+del {
			fx, err = o.f(x)
			if err != nil {
				return
			}
			sum += fx
		}
		o.s = 0.5 * (o.s + (o.b-o.a)*sum/float64(tnm))

		// This replaces s by its refined value.
		return o.s, nil
	}
}

// Integrate performs the numerical integration
func (o *ElementaryTrapz) Integrate() (res float64, err error) {
	jmax := 20
	var olds float64
	for j := 0; j < jmax; j++ {
		o.s, err = o.Next()
		if err != nil {
			return
		}
		if j > 5 {
			if math.Abs(o.s-olds) < o.eps*math.Abs(olds) || (o.s == 0 && olds == 0) {
				return o.s, nil
			}
		}
		olds = o.s
	}
	return 0, chk.Err("achieved maximum number of iterations (n=%d)", jmax)
}

// Simp structure implements the Simpson's method for quadrature with refinement.
type ElementarySimpson struct {
	n    int     // current level of refinement.
	a, b float64 // limits
	s    float64 // current value of the integral
	eps  float64 // precision
	f    fun.Ss  // the function
}

// Init initialises Simp structure
func (o *ElementarySimpson) Init(f fun.Ss, a, b, eps float64) {
	o.n = 0
	o.f = f
	o.a = a
	o.b = b
	o.eps = eps
}

// Next returns the nth stage of refinement of the extended trapezoidal rule. On the first call (n=1),
// R b the routine returns the crudest estimate of a f .x/dx. Subsequent calls set n=2,3,... and
// improve the accuracy by adding 2 n-2 additional interior points.
func (o *ElementarySimpson) Next() (res float64, err error) {
	var x, sum, del, fa, fb, fx float64
	var it, j, tnm int
	o.n += 1
	if o.n == 1 {
		fa, err = o.f(o.a)
		if err != nil {
			return
		}
		fb, err = o.f(o.b)
		if err != nil {
			return
		}
		o.s = 0.5 * (o.b - o.a) * (fa + fb)
		return o.s, nil
	} else {
		for it, j = 1, 1; j < o.n-1; j++ {
			it *= 2
		}
		tnm = it
		del = (o.b - o.a) / float64(tnm)

		// This is the spacing of the points to be added.
		x = o.a + 0.5*del

		for sum, j = 0.0, 0; j < it; j, x = j+1, x+del {
			fx, err = o.f(x)
			if err != nil {
				return
			}
			sum += fx
		}
		o.s = 0.5 * (o.s + (o.b-o.a)*sum/float64(tnm))

		// This replaces s by its refined value.
		return o.s, nil
	}
}

// Integrate performs the numerical integration
func (o *ElementarySimpson) Integrate() (res float64, err error) {
	jmax := 20
	var s, st, ost, os float64
	for j := 0; j < jmax; j++ {
		st, err = o.Next()
		if err != nil {
			return
		}
		s = (4*st - ost) / 3
		if j > 5 {
			if math.Abs(s-os) < o.eps*math.Abs(os) || (s == 0 && os == 0) {
				return s, nil
			}
		}
		os = s
		ost = st
	}
	return 0, chk.Err("achieved maximum number of iterations (n=%d)", jmax)
}
<commit_msg>Fix lint<commit_after>// Copyright 2016 The Gosl Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package num

import (
	"math"

	"github.com/cpmech/gosl/chk"
	"github.com/cpmech/gosl/fun"
)

// The algorithms below are based on [1]
// REFERENCES:
// [1] Press WH, Teukolsky SA, Vetterling WT, Fnannery BP (2007) Numerical Recipes: The Art of
//     Scientific Computing. Third Edition. Cambridge University Press. 1235p.

// QuadElementary defines the interface for elementary quadrature algorithms with refinement.
type QuadElementary interface {
	Init(f fun.Ss, a, b, eps float64) // The constructor takes as inputs f, the function or functor to be integrated between limits a and b, also input.
	Integrate() (float64, error)      // Returns the integral for the specified input data
}

// ElementaryTrapz structure is used for the trapezoidal integration rule with refinement.
type ElementaryTrapz struct {
	n    int     // current level of refinement.
	a, b float64 // limits
	s    float64 // current value of the integral
	eps  float64 // precision
	f    fun.Ss  // the function
}

// Init initialises Trap structure
func (o *ElementaryTrapz) Init(f fun.Ss, a, b, eps float64) {
	o.n = 0
	o.f = f
	o.a = a
	o.b = b
	o.eps = eps
}

// Next returns the nth stage of refinement of the extended trapezoidal rule. On the first call (n=1),
// R b the routine returns the crudest estimate of a f .x/dx. Subsequent calls set n=2,3,... and
// improve the accuracy by adding 2 n-2 additional interior points.
func (o *ElementaryTrapz) Next() (res float64, err error) {
	var x, sum, del float64
	var it, j, tnm int
	o.n++
	var fa, fb, fx float64
	if o.n == 1 {
		fa, err = o.f(o.a)
		if err != nil {
			return
		}
		fb, err = o.f(o.b)
		if err != nil {
			return
		}
		o.s = 0.5 * (o.b - o.a) * (fa + fb)
		return o.s, nil
	}
	for it, j = 1, 1; j < o.n-1; j++ {
		it *= 2
	}
	tnm = it
	del = (o.b - o.a) / float64(tnm)

	// spacing of the points to be added.
	x = o.a + 0.5*del

	for sum, j = 0.0, 0; j < it; j, x = j+1, x+del {
		fx, err = o.f(x)
		if err != nil {
			return
		}
		sum += fx
	}
	o.s = 0.5 * (o.s + (o.b-o.a)*sum/float64(tnm))

	// replace s by its refined value.
	return o.s, nil
}

// Integrate performs the numerical integration
func (o *ElementaryTrapz) Integrate() (res float64, err error) {
	jmax := 20
	var olds float64
	for j := 0; j < jmax; j++ {
		o.s, err = o.Next()
		if err != nil {
			return
		}
		if j > 5 {
			if math.Abs(o.s-olds) < o.eps*math.Abs(olds) || (o.s == 0 && olds == 0) {
				return o.s, nil
			}
		}
		olds = o.s
	}
	return 0, chk.Err("achieved maximum number of iterations (n=%d)", jmax)
}

// ElementarySimpson structure implements the Simpson's method for quadrature with refinement.
type ElementarySimpson struct {
	n    int     // current level of refinement.
	a, b float64 // limits
	s    float64 // current value of the integral
	eps  float64 // precision
	f    fun.Ss  // the function
}

// Init initialises Simp structure
func (o *ElementarySimpson) Init(f fun.Ss, a, b, eps float64) {
	o.n = 0
	o.f = f
	o.a = a
	o.b = b
	o.eps = eps
}

// Next returns the nth stage of refinement of the extended trapezoidal rule. On the first call (n=1),
// R b the routine returns the crudest estimate of a f .x/dx. Subsequent calls set n=2,3,... and
// improve the accuracy by adding 2 n-2 additional interior points.
func (o *ElementarySimpson) Next() (res float64, err error) {
	var x, sum, del, fa, fb, fx float64
	var it, j, tnm int
	o.n++
	if o.n == 1 {
		fa, err = o.f(o.a)
		if err != nil {
			return
		}
		fb, err = o.f(o.b)
		if err != nil {
			return
		}
		o.s = 0.5 * (o.b - o.a) * (fa + fb)
		return o.s, nil
	}
	for it, j = 1, 1; j < o.n-1; j++ {
		it *= 2
	}
	tnm = it
	del = (o.b - o.a) / float64(tnm)

	// spacing of the points to be added.
	x = o.a + 0.5*del

	for sum, j = 0.0, 0; j < it; j, x = j+1, x+del {
		fx, err = o.f(x)
		if err != nil {
			return
		}
		sum += fx
	}
	o.s = 0.5 * (o.s + (o.b-o.a)*sum/float64(tnm))

	// replace s by its refined value.
	return o.s, nil
}

// Integrate performs the numerical integration
func (o *ElementarySimpson) Integrate() (res float64, err error) {
	jmax := 20
	var s, st, ost, os float64
	for j := 0; j < jmax; j++ {
		st, err = o.Next()
		if err != nil {
			return
		}
		s = (4*st - ost) / 3
		if j > 5 {
			if math.Abs(s-os) < o.eps*math.Abs(os) || (s == 0 && os == 0) {
				return s, nil
			}
		}
		os = s
		ost = st
	}
	return 0, chk.Err("achieved maximum number of iterations (n=%d)", jmax)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	goetcd "github.com/coreos/go-etcd/etcd"
	"log"
	"strings"
	"time"
)

type Watch struct {
	client    *goetcd.Client
	prefix    string
	sentIndex uint64
	C         chan *goetcd.Node
}

func NewWatch(client *goetcd.Client, prefix string, limit int) *Watch {
	watch := &Watch{
		client: client,
		prefix: prefix,
		C:      make(chan *goetcd.Node, limit),
	}

	go watch.worker()

	return watch
}

func (watch *Watch) worker() {
	defer close(watch.C)

	if _, err := watch.client.SetDir(watch.prefix, 0); err != nil {
		// Ignore error code 102 (directory exists).
		if e, ok := err.(*goetcd.EtcdError); !ok || e.ErrorCode != 102 {
			log.Println("Watch etcd.SetDir error", watch.prefix, err)
			return
		}
	}

	for {
		// Fetch all current keys under this prefix, recursively.
		resp, err := watch.client.Get(watch.prefix, true, true)
		if err != nil {
			log.Println("Watch etcd.Get error", watch.prefix, err)
			return
		}

		// Start watching for updates after the current index given in the Get.
		index := resp.EtcdIndex

		log.Println("Watch etcd.Watch starting index", index, watch.prefix)

		// Find all non-directory nodes and send each to the channel.
		watch.sendNodes(resp.Node)

		for {
			// Fetch the next changed node for this prefix after index.
			resp, err = watch.client.Watch(watch.prefix, index+1, true, nil, nil)
			if err != nil {
				// TODO: Etcd closes the connection after 5 minutes,
				// resulting in a json.SyntaxError.  Retry watch.
				if _, ok := err.(*json.SyntaxError); ok {
					log.Println("Watch etcd.Watch retrying connection", watch.prefix)
					time.Sleep(*etcdRetryDelay)
					continue
				}

				// 401 means our index is too old, and we need to Get a new one.
				if e, ok := err.(*goetcd.EtcdError); ok && e.ErrorCode == 401 {
					log.Println("Watch etcd.Watch index", index+1, "too old", watch.prefix)
					time.Sleep(*etcdRetryDelay)
					break
				}

				log.Println("Watch etcd.Watch error", watch.prefix, err)
				return
			}

			// Send the changed node(s) to the update channel, track largest index we've sent.
			if i := watch.sendNodes(resp.Node); i > index {
				index = i
				watch.sentIndex = i
			}
		}
	}
}

func (watch *Watch) sendNodes(node *goetcd.Node) uint64 {
	if !node.Dir {
		// Send this to channel if it is not a repeat.
		if node.ModifiedIndex > watch.sentIndex {
			node.Key = strings.TrimPrefix(node.Key, watch.prefix+"/")
			watch.C <- node
		}
		return node.ModifiedIndex
	}

	var index uint64
	for _, node := range node.Nodes {
		// Iterate into directory, track largest index we've seen.
		if i := watch.sendNodes(node); i > index {
			index = i
		}
	}
	return index
}
<commit_msg>Also check for io.EOF for connection closing.<commit_after>package main

import (
	"encoding/json"
	goetcd "github.com/coreos/go-etcd/etcd"
	"io"
	"log"
	"strings"
	"time"
)

type Watch struct {
	client    *goetcd.Client
	prefix    string
	sentIndex uint64
	C         chan *goetcd.Node
}

func NewWatch(client *goetcd.Client, prefix string, limit int) *Watch {
	watch := &Watch{
		client: client,
		prefix: prefix,
		C:      make(chan *goetcd.Node, limit),
	}

	go watch.worker()

	return watch
}

func (watch *Watch) worker() {
	defer close(watch.C)

	if _, err := watch.client.SetDir(watch.prefix, 0); err != nil {
		// Ignore error code 102 (directory exists).
		if e, ok := err.(*goetcd.EtcdError); !ok || e.ErrorCode != 102 {
			log.Println("Watch etcd.SetDir error", watch.prefix, err)
			return
		}
	}

	for {
		// Fetch all current keys under this prefix, recursively.
		resp, err := watch.client.Get(watch.prefix, true, true)
		if err != nil {
			log.Println("Watch etcd.Get error", watch.prefix, err)
			return
		}

		// Start watching for updates after the current index given in the Get.
		index := resp.EtcdIndex

		log.Println("Watch etcd.Watch starting index", index, watch.prefix)

		// Find all non-directory nodes and send each to the channel.
		watch.sendNodes(resp.Node)

		for {
			// Fetch the next changed node for this prefix after index.
			resp, err = watch.client.Watch(watch.prefix, index+1, true, nil, nil)
			if err != nil {
				// TODO: Etcd closes the connection after 5 minutes,
				// resulting in a json.SyntaxError.  Retry watch.
				if _, ok := err.(*json.SyntaxError); ok || err == io.EOF {
					log.Println("Watch etcd.Watch retrying connection", watch.prefix)
					time.Sleep(*etcdRetryDelay)
					continue
				}

				// 401 means our index is too old, and we need to Get a new one.
				if e, ok := err.(*goetcd.EtcdError); ok && e.ErrorCode == 401 {
					log.Println("Watch etcd.Watch index", index+1, "too old", watch.prefix)
					time.Sleep(*etcdRetryDelay)
					break
				}

				log.Println("Watch etcd.Watch error", watch.prefix, err)
				return
			}

			// Send the changed node(s) to the update channel, track largest index we've sent.
			if i := watch.sendNodes(resp.Node); i > index {
				index = i
				watch.sentIndex = i
			}
		}
	}
}

func (watch *Watch) sendNodes(node *goetcd.Node) uint64 {
	if !node.Dir {
		// Send this to channel if it is not a repeat.
		if node.ModifiedIndex > watch.sentIndex {
			node.Key = strings.TrimPrefix(node.Key, watch.prefix+"/")
			watch.C <- node
		}
		return node.ModifiedIndex
	}

	var index uint64
	for _, node := range node.Nodes {
		// Iterate into directory, track largest index we've seen.
		if i := watch.sendNodes(node); i > index {
			index = i
		}
	}
	return index
}
<|endoftext|>