Merge pull request #40 from ruby/dbm-psych-4

Use `safe_load` instead of `load`

.github/workflows/test.yml

@@ -28,7 +28,12 @@ jobs:
       uses: ruby/setup-ruby@v1
       with:
         ruby-version: ${{ matrix.ruby }}
+    - name: Install system dependencies
+      run: |
+        sudo apt-get update
+        sudo apt-get install -y libdb-dev
+      if: matrix.os == 'ubuntu-latest'
     - name: Install dependencies
       run: bundle install --jobs 4 --retry 3
     - name: Run test
-      run: bundle exec rake test
+      run: rake test

Gemfile

@@ -1,8 +1,10 @@
 source "https://rubygems.org"
 
 group :development do
-  gem "bundler"
   gem "rake"
   gem "test-unit"
   gem "test-unit-ruby-core"
+
+  gem "pstore"
+  gem "dbm"
 end

lib/yaml/dbm.rb

@@ -1,6 +1,10 @@
 # frozen_string_literal: false
 require 'yaml'
-require 'dbm'
+
+begin
+  require 'dbm'
+rescue LoadError
+end
 
 module YAML
 
@@ -16,7 +20,6 @@ module YAML
 #
 # See the documentation for ::DBM and ::YAML for more information.
 class DBM < ::DBM
-    VERSION = "0.1" # :nodoc:
 
     # :call-seq:
     #   ydbm[key] -> value
@@ -56,7 +59,13 @@ def []=( key, val )
     def fetch( keystr, ifnone = nil )
         begin
             val = super( keystr )
-            return YAML.load( val ) if String === val
+            if String === val
+                if YAML.respond_to?(:safe_load)
+                    return YAML.safe_load( val )
+                else
+                    return YAML.load( val )
+                end
+            end
         rescue IndexError
         end
         if block_given?
@@ -102,7 +111,11 @@ def values_at( *keys )
     def delete( key )
         v = super( key )
         if String === v
-            v = YAML.load( v )
+            if YAML.respond_to?(:safe_load)
+                v = YAML.safe_load( v )
+            else
+                v = YAML.load( v )
+            end
         end
         v
     end
@@ -149,7 +162,7 @@ def each_pair # :yields: [key, value]
     #
     # Returns +self+.
     def each_value # :yields: value
-        super { |v| yield YAML.load( v ) }
+        super { |v| yield YAML.respond_to?(:safe_load) ? YAML.safe_load( v ) : YAML.load( v ) }
         self
     end
 
@@ -158,7 +171,7 @@ def each_value # :yields: value
     #
     # Returns an array of values from the database.
     def values
-        super.collect { |v| YAML.load( v ) }
+        super.collect { |v| YAML.respond_to?(:safe_load) ? YAML.safe_load( v ) : YAML.load( v ) }
     end
 
     # :call-seq:
@@ -204,7 +217,9 @@ def replace( hsh )
     # The order in which values are removed/returned is not guaranteed.
     def shift
         a = super
-        a[1] = YAML.load( a[1] ) if a
+        if a
+          a[1] = YAML.respond_to?(:safe_load) ? YAML.safe_load( a[1] ) : YAML.load( a[1] )
+        end
         a
     end
 
@@ -277,4 +292,4 @@ def to_hash
     alias :each :each_pair
 end
 
-end
+end if defined?(DBM)

test/yaml/test_dbm.rb

@@ -0,0 +1,46 @@
+require "test/unit"
+require "yaml/dbm"
+
+class TestYAMLDBM < Test::Unit::TestCase
+  def setup
+    @dbm = YAML::DBM.new("test")
+  end
+
+  def teardown
+    @dbm.close
+    File.unlink("test.db")
+  end
+
+  def test_fetch
+    @dbm["key"] = "value"
+    assert_equal "value", @dbm["key"]
+    assert_equal "value", @dbm.fetch("key")
+  end
+
+  def test_delete
+    @dbm["key"] = "value"
+    assert_equal "value", @dbm.delete("key")
+    assert_nil @dbm["key"]
+  end
+
+  def test_each_value
+    @dbm["key1"] = "value1"
+    @dbm["key2"] = "value2"
+    @dbm.each_value do |value|
+      assert_match(/value[12]/, value)
+    end
+  end
+
+  def test_values
+    @dbm["key1"] = "value1"
+    @dbm["key2"] = "value2"
+    @dbm.values.each do |value|
+      assert_match(/value[12]/, value)
+    end
+  end
+
+  def test_shift
+    @dbm["key"] = "value"
+    assert_equal ["key", "value"], @dbm.shift
+  end
+end if defined?(YAML::DBM)