Add FatalHook for graceful shutdown in case of Fatal log

This commit is contained in:
Elara 2020-12-03 09:32:03 -08:00
parent 5f2e020e75
commit 73b717f2e3
8 changed files with 40 additions and 19 deletions

View File

@ -26,7 +26,7 @@ func NewConfig(actionType string, actionData string) *Config {
// Create config file // Create config file
func (config *Config) CreateFile(dir string) { func (config *Config) CreateFile(dir string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Create config file at given directory // Create config file at given directory
configFile, err := os.Create(dir + "/config.json") configFile, err := os.Create(dir + "/config.json")
if err != nil { log.Fatal().Err(err).Msg("Error creating config file") } if err != nil { log.Fatal().Err(err).Msg("Error creating config file") }
@ -45,7 +45,7 @@ func (config *Config) CreateFile(dir string) {
// Collect all required files into given directory // Collect all required files into given directory
func (config *Config) CollectFiles(dir string) { func (config *Config) CollectFiles(dir string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// If action type is file // If action type is file
if config.ActionType == "file" { if config.ActionType == "file" {
// Open file path in config.ActionData // Open file path in config.ActionData
@ -69,7 +69,7 @@ func (config *Config) CollectFiles(dir string) {
// Read config file at given file path // Read config file at given file path
func (config *Config) ReadFile(filePath string) { func (config *Config) ReadFile(filePath string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Read file at filePath // Read file at filePath
fileData, err := ioutil.ReadFile(filePath) fileData, err := ioutil.ReadFile(filePath)
if err != nil { log.Fatal().Err(err).Msg("Error reading config file") } if err != nil { log.Fatal().Err(err).Msg("Error reading config file") }
@ -81,7 +81,7 @@ func (config *Config) ReadFile(filePath string) {
// Execute action specified in config // Execute action specified in config
func (config *Config) ExecuteAction(srcDir string) { func (config *Config) ExecuteAction(srcDir string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// If action is file // If action is file
if config.ActionType == "file" { if config.ActionType == "file" {
// Open file from config at given directory // Open file from config at given directory

View File

@ -12,7 +12,7 @@ import (
// Discover opensend receivers on the network // Discover opensend receivers on the network
func DiscoverReceivers() ([]string, []string) { func DiscoverReceivers() ([]string, []string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Create zeroconf resolver // Create zeroconf resolver
resolver, err := zeroconf.NewResolver(nil) resolver, err := zeroconf.NewResolver(nil)
if err != nil { log.Fatal().Err(err).Msg("Error creating zeroconf resolver") } if err != nil { log.Fatal().Err(err).Msg("Error creating zeroconf resolver") }

View File

@ -19,7 +19,7 @@ import (
// Encrypt given file using the shared key // Encrypt given file using the shared key
func EncryptFile(filePath string, newFilePath string, sharedKey string) { func EncryptFile(filePath string, newFilePath string, sharedKey string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Read data from file // Read data from file
data, err := ioutil.ReadFile(filePath) data, err := ioutil.ReadFile(filePath)
if err != nil { log.Fatal().Err(err).Msg("Error reading file") } if err != nil { log.Fatal().Err(err).Msg("Error reading file") }
@ -54,7 +54,7 @@ func EncryptFile(filePath string, newFilePath string, sharedKey string) {
// Decrypt given file using the shared key // Decrypt given file using the shared key
func DecryptFile(filePath string, newFilePath string, sharedKey string) { func DecryptFile(filePath string, newFilePath string, sharedKey string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Read data from file // Read data from file
data, err := ioutil.ReadFile(filePath) data, err := ioutil.ReadFile(filePath)
if err != nil { log.Fatal().Err(err).Msg("Error reading file") } if err != nil { log.Fatal().Err(err).Msg("Error reading file") }
@ -88,7 +88,7 @@ func DecryptFile(filePath string, newFilePath string, sharedKey string) {
// Encrypt files in given directory using shared key // Encrypt files in given directory using shared key
func EncryptFiles(dir string, sharedKey string) { func EncryptFiles(dir string, sharedKey string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Walk given directory // Walk given directory
err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error { err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
// If error reading, return err // If error reading, return err
@ -110,7 +110,7 @@ func EncryptFiles(dir string, sharedKey string) {
// Decrypt files in given directory using shared key // Decrypt files in given directory using shared key
func DecryptFiles(dir string, sharedKey string) { func DecryptFiles(dir string, sharedKey string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Walk given directory // Walk given directory
err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error { err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
// If error reading, return err // If error reading, return err

View File

@ -18,7 +18,7 @@ import (
// Save encrypted key to file // Save encrypted key to file
func SaveEncryptedKey(encryptedKey []byte, filePath string) { func SaveEncryptedKey(encryptedKey []byte, filePath string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Create file at given file path // Create file at given file path
keyFile, err := os.Create(filePath) keyFile, err := os.Create(filePath)
if err != nil { log.Fatal().Err(err).Msg("Error creating file") } if err != nil { log.Fatal().Err(err).Msg("Error creating file") }
@ -34,7 +34,7 @@ func SaveEncryptedKey(encryptedKey []byte, filePath string) {
// Create HTTP server to transmit files // Create HTTP server to transmit files
func SendFiles(dir string) { func SendFiles(dir string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Instantiate http.Server struct // Instantiate http.Server struct
srv := &http.Server{} srv := &http.Server{}
// Listen on all ipv4 addresses on port 9898 // Listen on all ipv4 addresses on port 9898
@ -114,7 +114,7 @@ func SendFiles(dir string) {
// Get files from sender // Get files from sender
func RecvFiles(dir string, senderAddr string) { func RecvFiles(dir string, senderAddr string) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Get server address by getting the IP without the port, prepending http:// and appending :9898 // Get server address by getting the IP without the port, prepending http:// and appending :9898
serverAddr := "http://" + strings.Split(senderAddr, ":")[0] + ":9898" serverAddr := "http://" + strings.Split(senderAddr, ":")[0] + ":9898"
// GET /index on sender's HTTP server // GET /index on sender's HTTP server

View File

@ -15,7 +15,7 @@ import (
// Generate RSA keypair // Generate RSA keypair
func GenerateRSAKeypair() (*rsa.PrivateKey, *rsa.PublicKey) { func GenerateRSAKeypair() (*rsa.PrivateKey, *rsa.PublicKey) {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Generate private/public RSA keypair // Generate private/public RSA keypair
privateKey, err := rsa.GenerateKey(rand.Reader, 2048) privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil { log.Fatal().Err(err).Msg("Error generating RSA keypair") } if err != nil { log.Fatal().Err(err).Msg("Error generating RSA keypair") }
@ -28,7 +28,7 @@ func GenerateRSAKeypair() (*rsa.PrivateKey, *rsa.PublicKey) {
// Get public key from sender // Get public key from sender
func GetKey(senderAddr string) []byte { func GetKey(senderAddr string) []byte {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Get server address by getting the IP without the port, prepending http:// and appending :9898 // Get server address by getting the IP without the port, prepending http:// and appending :9898
serverAddr := "http://" + strings.Split(senderAddr, ":")[0] + ":9898" serverAddr := "http://" + strings.Split(senderAddr, ":")[0] + ":9898"
// GET /key on the sender's HTTP server // GET /key on the sender's HTTP server
@ -55,7 +55,7 @@ func GetKey(senderAddr string) []byte {
// Encrypt shared key with received public key // Encrypt shared key with received public key
func EncryptKey(sharedKey string, recvPubKey *rsa.PublicKey) []byte { func EncryptKey(sharedKey string, recvPubKey *rsa.PublicKey) []byte {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Encrypt shared key using RSA // Encrypt shared key using RSA
encryptedSharedKey, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, recvPubKey, []byte(sharedKey), nil) encryptedSharedKey, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, recvPubKey, []byte(sharedKey), nil)
if err != nil { log.Fatal().Err(err).Msg("Error encrypting shared key") } if err != nil { log.Fatal().Err(err).Msg("Error encrypting shared key") }

View File

@ -12,7 +12,7 @@ import (
// Exchange keys with sender // Exchange keys with sender
func ReceiverKeyExchange(key *rsa.PublicKey) string { func ReceiverKeyExchange(key *rsa.PublicKey) string {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Create TCP listener on port 9797 // Create TCP listener on port 9797
listener, err := net.Listen("tcp", ":9797") listener, err := net.Listen("tcp", ":9797")
if err != nil { log.Fatal().Err(err).Msg("Error starting listener") } if err != nil { log.Fatal().Err(err).Msg("Error starting listener") }
@ -49,7 +49,7 @@ func ReceiverKeyExchange(key *rsa.PublicKey) string {
// Exchange keys with receiver // Exchange keys with receiver
func SenderKeyExchange(receiverIP string) *rsa.PublicKey { func SenderKeyExchange(receiverIP string) *rsa.PublicKey {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
// Connect to TCP socket on receiver IP port 9797 // Connect to TCP socket on receiver IP port 9797
connection, err := net.Dial("tcp", receiverIP + ":9797") connection, err := net.Dial("tcp", receiverIP + ":9797")
if err != nil { log.Fatal().Err(err).Msg("Error connecting to sender") } if err != nil { log.Fatal().Err(err).Msg("Error connecting to sender") }

16
logging.go Normal file
View File

@ -0,0 +1,16 @@
package main
import (
"github.com/rs/zerolog"
"os"
)
type FatalHook struct {}
func (hook FatalHook) Run(_ *zerolog.Event, level zerolog.Level, _ string) {
// If log event is fatal
if level == zerolog.FatalLevel {
// Attempt removal of opensend directory
_ = os.RemoveAll(opensendDir)
}
}

View File

@ -17,15 +17,20 @@ import (
"time" "time"
) )
var opensendDir string
func main() { func main() {
// Use ConsoleWriter logger // Use ConsoleWriter logger
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}).Hook(FatalHook{})
if os.Args[1] == "f" {
log.Fatal().Msg("Test")
}
// Get user's home directory // Get user's home directory
homeDir, err := os.UserHomeDir() homeDir, err := os.UserHomeDir()
if err != nil { log.Fatal().Err(err).Msg("Error getting home directory") } if err != nil { log.Fatal().Err(err).Msg("Error getting home directory") }
// Define opensend directory as ~/.opensend // Define opensend directory as ~/.opensend
opensendDir := homeDir + "/.opensend" opensendDir = homeDir + "/.opensend"
// Create channel for signals // Create channel for signals
sig := make(chan os.Signal, 1) sig := make(chan os.Signal, 1)