From 51ae45a7995db54356f6975499e4624617442f9b Mon Sep 17 00:00:00 2001
From: ysqi <devysq@gmail.com>
Date: Sat, 23 Jan 2016 14:53:52 +0800
Subject: [PATCH] Fixed #1607

---
 config/ini.go      | 40 +++++++++++++++++++++++++++--------
 config/ini_test.go | 52 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 9 deletions(-)

diff --git a/config/ini.go b/config/ini.go
index 59e84e1e..5a700a0b 100644
--- a/config/ini.go
+++ b/config/ini.go
@@ -27,7 +27,6 @@ import (
 	"strings"
 	"sync"
 	"time"
-	"unicode"
 )
 
 var (
@@ -97,9 +96,11 @@ func (ini *IniConfig) parseFile(name string) (*IniConfigContainer, error) {
 		}
 		if bComment != nil {
 			line = bytes.TrimLeft(line, string(bComment))
-			line = bytes.TrimLeftFunc(line, unicode.IsSpace)
+			// Need append to a new line if multi-line comments.
+			if comment.Len() > 0 {
+				comment.WriteByte('\n')
+			}
 			comment.Write(line)
-			comment.WriteByte('\n')
 			continue
 		}
 
@@ -299,14 +300,35 @@ func (c *IniConfigContainer) SaveConfigFile(filename string) (err error) {
 	}
 	defer f.Close()
 
+	// Get section or key comments. Fixed #1607
+	getCommentStr := func(section, key string) string {
+		comment, ok := "", false
+		if len(key) == 0 {
+			comment, ok = c.sectionComment[section]
+		} else {
+			comment, ok = c.keyComment[section+"."+key]
+		}
+
+		if ok {
+			// Empty comment
+			if len(comment) == 0 || len(strings.TrimSpace(comment)) == 0 {
+				return string(bNumComment)
+			}
+			prefix := string(bNumComment)
+			// Add the line head character "#"
+			return prefix + strings.Replace(comment, lineBreak, lineBreak+prefix, -1)
+		}
+		return ""
+	}
+
 	buf := bytes.NewBuffer(nil)
 	// Save default section at first place
 	if dt, ok := c.data[defaultSection]; ok {
 		for key, val := range dt {
 			if key != " " {
 				// Write key comments.
-				if v, ok := c.keyComment[key]; ok {
-					if _, err = buf.WriteString(string(bNumComment) + v + lineBreak); err != nil {
+				if v := getCommentStr(defaultSection, key); len(v) > 0 {
+					if _, err = buf.WriteString(v + lineBreak); err != nil {
 						return err
 					}
 				}
@@ -327,8 +349,8 @@ func (c *IniConfigContainer) SaveConfigFile(filename string) (err error) {
 	for section, dt := range c.data {
 		if section != defaultSection {
 			// Write section comments.
-			if v, ok := c.sectionComment[section]; ok {
-				if _, err = buf.WriteString(string(bNumComment) + v + lineBreak); err != nil {
+			if v := getCommentStr(section, ""); len(v) > 0 {
+				if _, err = buf.WriteString(v + lineBreak); err != nil {
 					return err
 				}
 			}
@@ -341,8 +363,8 @@ func (c *IniConfigContainer) SaveConfigFile(filename string) (err error) {
 			for key, val := range dt {
 				if key != " " {
 					// Write key comments.
-					if v, ok := c.keyComment[key]; ok {
-						if _, err = buf.WriteString(string(bNumComment) + v + lineBreak); err != nil {
+					if v := getCommentStr(section, key); len(v) > 0 {
+						if _, err = buf.WriteString(v + lineBreak); err != nil {
 							return err
 						}
 					}
diff --git a/config/ini_test.go b/config/ini_test.go
index 7599ab8b..9f0bb8af 100644
--- a/config/ini_test.go
+++ b/config/ini_test.go
@@ -15,6 +15,7 @@
 package config
 
 import (
+	"io/ioutil"
 	"os"
 	"testing"
 )
@@ -101,3 +102,54 @@ func TestIni(t *testing.T) {
 	}
 
 }
+
+func TestIniSave(t *testing.T) {
+
+	const (
+		inicontext = `
+app = app
+;comment one
+#comment two
+# comment three
+appname = beeapi
+httpport = 8080
+# DB Info
+# enable db
+[dbinfo]
+# db type name
+# suport mysql,sqlserver
+name = mysql
+`
+
+		saveResult = `app=app
+#comment one
+#comment two
+# comment three
+appname=beeapi
+httpport=8080
+
+# DB Info
+# enable db
+[dbinfo]
+# db type name
+# suport mysql,sqlserver
+name=mysql
+
+`
+	)
+	cfg, err := NewConfigData("ini", []byte(inicontext))
+	if err != nil {
+		t.Fatal(err)
+	}
+	name := "newIniConfig.ini"
+	if err := cfg.SaveConfigFile(name); err != nil {
+		t.Fatal(err)
+	}
+	defer os.Remove(name)
+
+	if data, err := ioutil.ReadFile(name); err != nil {
+		t.Fatal(err)
+	} else if string(data) != saveResult {
+		t.Fatal("different after save ini config file.")
+	}
+}