From e26d6aedc15334f7c2bca88e37037f717659d9b8 Mon Sep 17 00:00:00 2001
From: Yauhen Lazurkin <foryauhen@gmail.com>
Date: Sat, 24 Sep 2016 02:20:44 +0300
Subject: [PATCH] Support time.Duration in viper.Unmarshal (#205)

* Fixes #105
---
 viper.go      | 20 +++++++++++++-------
 viper_test.go | 10 ++++++----
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/viper.go b/viper.go
index 8a32482..fafabde 100644
--- a/viper.go
+++ b/viper.go
@@ -624,7 +624,7 @@ func (v *Viper) UnmarshalKey(key string, rawVal interface{}) error {
 // on the fields of the structure are properly set.
 func Unmarshal(rawVal interface{}) error { return v.Unmarshal(rawVal) }
 func (v *Viper) Unmarshal(rawVal interface{}) error {
-	err := mapstructure.WeakDecode(v.AllSettings(), rawVal)
+	err := decode(v.AllSettings(), defaultDecoderConfig(rawVal))
 
 	if err != nil {
 		return err
@@ -635,16 +635,19 @@ func (v *Viper) Unmarshal(rawVal interface{}) error {
 	return nil
 }
 
-// A wrapper around mapstructure.Decode that mimics the WeakDecode functionality
-// while erroring on non existing vals in the destination struct.
-func weakDecodeExact(input, output interface{}) error {
-	config := &mapstructure.DecoderConfig{
-		ErrorUnused:      true,
+// defaultDecoderConfig returns default mapsstructure.DecoderConfig with suppot
+// of time.Duration values
+func defaultDecoderConfig(output interface{}) *mapstructure.DecoderConfig {
+	return &mapstructure.DecoderConfig{
 		Metadata:         nil,
 		Result:           output,
 		WeaklyTypedInput: true,
+		DecodeHook:       mapstructure.StringToTimeDurationHookFunc(),
 	}
+}
 
+// A wrapper around mapstructure.Decode that mimics the WeakDecode functionality
+func decode(input interface{}, config *mapstructure.DecoderConfig) error {
 	decoder, err := mapstructure.NewDecoder(config)
 	if err != nil {
 		return err
@@ -655,7 +658,10 @@ func weakDecodeExact(input, output interface{}) error {
 // UnmarshalExact unmarshals the config into a Struct, erroring if a field is nonexistent
 // in the destination struct.
 func (v *Viper) UnmarshalExact(rawVal interface{}) error {
-	err := weakDecodeExact(v.AllSettings(), rawVal)
+	config := defaultDecoderConfig(rawVal)
+	config.ErrorUnused = true
+
+	err := decode(v.AllSettings(), config)
 
 	if err != nil {
 		return err
diff --git a/viper_test.go b/viper_test.go
index cf897d6..127e650 100644
--- a/viper_test.go
+++ b/viper_test.go
@@ -453,10 +453,12 @@ func TestRecursiveAliases(t *testing.T) {
 func TestUnmarshal(t *testing.T) {
 	SetDefault("port", 1313)
 	Set("name", "Steve")
+	Set("duration", "1s1ms")
 
 	type config struct {
-		Port int
-		Name string
+		Port     int
+		Name     string
+		Duration time.Duration
 	}
 
 	var C config
@@ -466,14 +468,14 @@ func TestUnmarshal(t *testing.T) {
 		t.Fatalf("unable to decode into struct, %v", err)
 	}
 
-	assert.Equal(t, &C, &config{Name: "Steve", Port: 1313})
+	assert.Equal(t, &C, &config{Name: "Steve", Port: 1313, Duration: time.Second + time.Millisecond})
 
 	Set("port", 1234)
 	err = Unmarshal(&C)
 	if err != nil {
 		t.Fatalf("unable to decode into struct, %v", err)
 	}
-	assert.Equal(t, &C, &config{Name: "Steve", Port: 1234})
+	assert.Equal(t, &C, &config{Name: "Steve", Port: 1234, Duration: time.Second + time.Millisecond})
 }
 
 func TestBindPFlags(t *testing.T) {