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

// Package pemfile provides a file watching certificate provider plugin
// implementation which works for files with PEM contents.
//
// Experimental
//
// Notice: All APIs in this package are experimental and may be removed in a
// later release.
package pemfile

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"time"

	"google.golang.org/grpc/credentials/tls/certprovider"
	"google.golang.org/grpc/grpclog"
)

const (
	defaultCertRefreshDuration = 1 * time.Hour
	defaultRootRefreshDuration = 2 * time.Hour
)

var (
	// For overriding from unit tests.
	newDistributor = func() distributor { return certprovider.NewDistributor() }

	logger = grpclog.Component("pemfile")
)

// Options configures a certificate provider plugin that watches a specified set
// of files that contain certificates and keys in PEM format.
type Options struct {
	// CertFile is the file that holds the identity certificate.
	// Optional. If this is set, KeyFile must also be set.
	CertFile string
	// KeyFile is the file that holds identity private key.
	// Optional. If this is set, CertFile must also be set.
	KeyFile string
	// RootFile is the file that holds trusted root certificate(s).
	// Optional.
	RootFile string
	// CertRefreshDuration is the amount of time the plugin waits before
	// checking for updates in the specified identity certificate and key file.
	// Optional. If not set, a default value (1 hour) will be used.
	CertRefreshDuration time.Duration
	// RootRefreshDuration is the amount of time the plugin waits before
	// checking for updates in the specified root file.
	// Optional. If not set, a default value (2 hour) will be used.
	RootRefreshDuration time.Duration
}

// NewProvider returns a new certificate provider plugin that is configured to
// watch the PEM files specified in the passed in options.
func NewProvider(o Options) (certprovider.Provider, error) {
	if o.CertFile == "" && o.KeyFile == "" && o.RootFile == "" {
		return nil, fmt.Errorf("pemfile: at least one credential file needs to be specified")
	}
	if keySpecified, certSpecified := o.KeyFile != "", o.CertFile != ""; keySpecified != certSpecified {
		return nil, fmt.Errorf("pemfile: private key file and identity cert file should be both specified or not specified")
	}
	if o.CertRefreshDuration == 0 {
		o.CertRefreshDuration = defaultCertRefreshDuration
	}
	if o.RootRefreshDuration == 0 {
		o.RootRefreshDuration = defaultRootRefreshDuration
	}

	provider := &watcher{opts: o}
	if o.CertFile != "" && o.KeyFile != "" {
		provider.identityDistributor = newDistributor()
	}
	if o.RootFile != "" {
		provider.rootDistributor = newDistributor()
	}

	ctx, cancel := context.WithCancel(context.Background())
	provider.cancel = cancel
	go provider.run(ctx)

	return provider, nil
}

// watcher is a certificate provider plugin that implements the
// certprovider.Provider interface. It watches a set of certificate and key
// files and provides the most up-to-date key material for consumption by
// credentials implementation.
type watcher struct {
	identityDistributor distributor
	rootDistributor     distributor
	opts                Options
	certFileContents    []byte
	keyFileContents     []byte
	rootFileContents    []byte
	cancel              context.CancelFunc
}

// distributor wraps the methods on certprovider.Distributor which are used by
// the plugin. This is very useful in tests which need to know exactly when the
// plugin updates its key material.
type distributor interface {
	KeyMaterial(ctx context.Context) (*certprovider.KeyMaterial, error)
	Set(km *certprovider.KeyMaterial, err error)
	Stop()
}

// updateIdentityDistributor checks if the cert/key files that the plugin is
// watching have changed, and if so, reads the new contents and updates the
// identityDistributor with the new key material.
//
// Skips updates when file reading or parsing fails.
// TODO(easwars): Retry with limit (on the number of retries or the amount of
// time) upon failures.
func (w *watcher) updateIdentityDistributor() {
	if w.identityDistributor == nil {
		return
	}

	certFileContents, err := ioutil.ReadFile(w.opts.CertFile)
	if err != nil {
		logger.Warningf("certFile (%s) read failed: %v", w.opts.CertFile, err)
		return
	}
	keyFileContents, err := ioutil.ReadFile(w.opts.KeyFile)
	if err != nil {
		logger.Warningf("keyFile (%s) read failed: %v", w.opts.KeyFile, err)
		return
	}
	// If the file contents have not changed, skip updating the distributor.
	if bytes.Equal(w.certFileContents, certFileContents) && bytes.Equal(w.keyFileContents, keyFileContents) {
		return
	}

	cert, err := tls.X509KeyPair(certFileContents, keyFileContents)
	if err != nil {
		logger.Warningf("tls.X509KeyPair(%q, %q) failed: %v", certFileContents, keyFileContents, err)
		return
	}
	w.certFileContents = certFileContents
	w.keyFileContents = keyFileContents
	w.identityDistributor.Set(&certprovider.KeyMaterial{Certs: []tls.Certificate{cert}}, nil)
}

// updateRootDistributor checks if the root cert file that the plugin is
// watching hs changed, and if so, updates the rootDistributor with the new key
// material.
//
// Skips updates when root cert reading or parsing fails.
// TODO(easwars): Retry with limit (on the number of retries or the amount of
// time) upon failures.
func (w *watcher) updateRootDistributor() {
	if w.rootDistributor == nil {
		return
	}

	rootFileContents, err := ioutil.ReadFile(w.opts.RootFile)
	if err != nil {
		logger.Warningf("rootFile (%s) read failed: %v", w.opts.RootFile, err)
		return
	}
	trustPool := x509.NewCertPool()
	if !trustPool.AppendCertsFromPEM(rootFileContents) {
		logger.Warning("failed to parse root certificate")
		return
	}
	// If the file contents have not changed, skip updating the distributor.
	if bytes.Equal(w.rootFileContents, rootFileContents) {
		return
	}

	w.rootFileContents = rootFileContents
	w.rootDistributor.Set(&certprovider.KeyMaterial{Roots: trustPool}, nil)
}

// run is a long running goroutine which watches the configured files for
// changes, and pushes new key material into the appropriate distributors which
// is returned from calls to KeyMaterial().
func (w *watcher) run(ctx context.Context) {
	// Update both root and identity certs at the beginning. Subsequently,
	// update only the appropriate file whose ticker has fired.
	w.updateIdentityDistributor()
	w.updateRootDistributor()

	identityTicker := time.NewTicker(w.opts.CertRefreshDuration)
	rootTicker := time.NewTicker(w.opts.RootRefreshDuration)
	for {
		select {
		case <-ctx.Done():
			identityTicker.Stop()
			rootTicker.Stop()
			if w.identityDistributor != nil {
				w.identityDistributor.Stop()
			}
			if w.rootDistributor != nil {
				w.rootDistributor.Stop()
			}
			return
		case <-identityTicker.C:
			w.updateIdentityDistributor()
		case <-rootTicker.C:
			w.updateRootDistributor()
		}
	}
}

// KeyMaterial returns the key material sourced by the watcher.
// Callers are expected to use the returned value as read-only.
func (w *watcher) KeyMaterial(ctx context.Context) (*certprovider.KeyMaterial, error) {
	km := &certprovider.KeyMaterial{}
	if w.identityDistributor != nil {
		identityKM, err := w.identityDistributor.KeyMaterial(ctx)
		if err != nil {
			return nil, err
		}
		km.Certs = identityKM.Certs
	}
	if w.rootDistributor != nil {
		rootKM, err := w.rootDistributor.KeyMaterial(ctx)
		if err != nil {
			return nil, err
		}
		km.Roots = rootKM.Roots
	}
	return km, nil
}

// Close cleans up resources allocated by the watcher.
func (w *watcher) Close() {
	w.cancel()
}